//! xmodel-harness — drives the two-cohort cross-model authoring-form test.

use anyhow::{anyhow, bail, Context, Result};
use clap::Parser;
use std::collections::BTreeSet;
use std::path::PathBuf;
use xmodel_harness::ionos::{ChatRequest, ChatResponse, IonosClient, IonosError, Message};
use xmodel_harness::mock::MockResponses;
use xmodel_harness::pipeline::{preflight, run_pipeline, Cohort};
use xmodel_harness::scoring::{write_scores_csv, write_summary_md, FinalStatus, ScoreRow};
use xmodel_harness::strip_locations::strip_locations;
use xmodel_harness::tasks::{self, Task};

#[derive(Parser, Debug)]
#[command(name = "xmodel-harness", version, about = "Two-cohort cross-model authoring-form test")]
struct Args {
    #[arg(long)] rendered: PathBuf,
    #[arg(long)] tasks: PathBuf,
    #[arg(long)] out: PathBuf,
    #[arg(long)] model: String,
    #[arg(long, default_value_t = 5)] max_turns: u32,
    #[arg(long, default_value_t = 500_000)] token_budget: u64,
    #[arg(long)] mock: Option<PathBuf>,
}

enum Backend {
    Live(IonosClient),
    Mock(MockResponses),
}

fn main() -> Result<()> {
    let args = Args::parse();
    let run_dir = create_run_dir(&args.out)?;

    preflight().context("preflight: ail and clang must be on PATH")?;

    let backend = if let Some(mp) = &args.mock {
        Backend::Mock(MockResponses::load(mp)?)
    } else {
        Backend::Live(IonosClient::from_env().context("IONOS_API_TOKEN")?)
    };

    let rendered_json = std::fs::read_to_string(args.rendered.join("json.md"))
        .with_context(|| format!("reading {}", args.rendered.join("json.md").display()))?;
    let rendered_ail = std::fs::read_to_string(args.rendered.join("ail.md"))
        .with_context(|| format!("reading {}", args.rendered.join("ail.md").display()))?;
    let tasks = tasks::load_all(&args.tasks)?;
    if tasks.is_empty() { bail!("no tasks found in {}", args.tasks.display()); }

    let mut tokens_used: u64 = 0;
    let mut rows: Vec<ScoreRow> = Vec::new();
    let mut run_status = "ok";

    let mut completed: std::collections::BTreeSet<(&'static str, String)> = Default::default();
    'outer: for cohort in [Cohort::Json, Cohort::Ail] {
        let system_prompt = match cohort { Cohort::Json => &rendered_json, Cohort::Ail => &rendered_ail };
        for t in &tasks {
            let (row, consumed) = run_one(
                &backend, &args, cohort, system_prompt, t,
                &run_dir, args.token_budget.saturating_sub(tokens_used),
            )?;
            tokens_used = tokens_used.saturating_add(consumed);
            let aborted_budget = matches!(row.final_status, FinalStatus::BudgetAbort);
            completed.insert((cohort.as_str(), t.id.clone()));
            rows.push(row);
            if tokens_used >= args.token_budget {
                run_status = "budget_exceeded";
                break 'outer;
            }
            let _ = aborted_budget; // continue to next task regardless
        }
    }

    // If the budget ran out mid-run, the outer loop broke before
    // walking every (cohort, task) pair. The pairs we never reached
    // still need a row, with final_status = budget_abort.
    if run_status == "budget_exceeded" {
        for cohort in [Cohort::Json, Cohort::Ail] {
            for t in &tasks {
                if !completed.contains(&(cohort.as_str(), t.id.clone())) {
                    rows.push(ScoreRow {
                        cohort: cohort.as_str().to_string(),
                        task_id: t.id.clone(),
                        first_attempt_green: false,
                        turns_to_green: None,
                        prompt_tokens: 0,
                        completion_tokens: 0,
                        error_classes: BTreeSet::new(),
                        final_status: FinalStatus::BudgetAbort,
                    });
                }
            }
        }
    }

    write_scores_csv(&rows, &run_dir.join("scores.csv"))?;
    write_summary_md(&rows, &run_dir.join("summary.md"))?;
    std::fs::write(run_dir.join("RUN_STATUS"), run_status)?;
    eprintln!("xmodel-harness: run complete at {}; status={run_status}", run_dir.display());
    Ok(())
}

fn create_run_dir(out: &std::path::Path) -> Result<PathBuf> {
    let date = chrono::Utc::now().format("%Y-%m-%d");
    let hash: String = (0..6).map(|_| {
        let n = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().subsec_nanos();
        std::char::from_digit((n % 16) as u32, 16).unwrap()
    }).collect();
    let dir = out.join(format!("{date}-{hash}"));
    std::fs::create_dir_all(&dir)?;
    Ok(dir)
}

fn call_backend(backend: &Backend, model: &str, messages: &[Message], cohort: &str, task: &str, turn: u32)
    -> std::result::Result<ChatResponse, IonosError>
{
    match backend {
        Backend::Live(c) => c.post(&ChatRequest { model, messages, temperature: 0.0, top_p: 1.0 }),
        Backend::Mock(m) => m.response_for(cohort, task, turn).map_err(|e| IonosError::Transport(e.to_string())),
    }
}

fn run_one(
    backend: &Backend, args: &Args, cohort: Cohort, system_prompt: &str, task: &Task,
    run_dir: &std::path::Path, budget_remaining: u64,
) -> Result<(ScoreRow, u64)> {
    let cohort_dir = run_dir.join("per_cohort").join(cohort.as_str()).join(&task.id);
    std::fs::create_dir_all(&cohort_dir)?;
    let mut messages = vec![
        Message { role: "system".into(), content: system_prompt.to_string() },
        Message { role: "user".into(),   content: task.description.clone() },
    ];
    let mut prompt_total: u64 = 0;
    let mut completion_total: u64 = 0;
    let mut error_classes: BTreeSet<String> = BTreeSet::new();
    let mut first_attempt_green = false;
    let mut turns_to_green: Option<u32> = None;
    let mut final_status = FinalStatus::TurnLimit;

    for turn in 1..=args.max_turns {
        if prompt_total.saturating_add(completion_total) >= budget_remaining {
            final_status = FinalStatus::BudgetAbort;
            break;
        }
        let resp = match call_backend(backend, &args.model, &messages, cohort.as_str(), &task.id, turn) {
            Ok(r) => r,
            Err(IonosError::Auth(_)) | Err(IonosError::Transport(_)) | Err(IonosError::RetriesExhausted{..}) => {
                final_status = FinalStatus::ApiFailure;
                break;
            }
        };
        prompt_total += resp.usage.prompt_tokens;
        completion_total += resp.usage.completion_tokens;
        let program = resp.choices.first()
            .ok_or_else(|| anyhow!("empty choices array"))?
            .message.content.clone();

        std::fs::write(
            cohort_dir.join(format!("turn_{turn}_program.{}", cohort.extension())),
            &program,
        )?;
        std::fs::write(
            cohort_dir.join(format!("turn_{turn}_request.json")),
            serde_json::to_string_pretty(&messages)?,
        )?;
        std::fs::write(
            cohort_dir.join(format!("turn_{turn}_response.json")),
            serde_json::to_string_pretty(&serde_json::json!({
                "usage": { "prompt_tokens": resp.usage.prompt_tokens, "completion_tokens": resp.usage.completion_tokens, "total_tokens": resp.usage.total_tokens },
                "content": program,
            }))?,
        )?;

        let workdir = tempfile::Builder::new().prefix("cma2-run-").tempdir()?;
        let cap = run_pipeline(cohort, &program, &task.expected_stdout, workdir.path())?;
        std::fs::write(cohort_dir.join(format!("turn_{turn}_check_stderr.txt")), &cap.stderr)?;
        std::fs::write(cohort_dir.join(format!("turn_{turn}_run_stdout.txt")), &cap.stdout)?;

        match cap.error {
            None => {
                if turn == 1 { first_attempt_green = true; }
                turns_to_green = Some(turn);
                final_status = FinalStatus::Green;
                break;
            }
            Some(err) => {
                if let Some((class, _)) = err.split_once(':') {
                    error_classes.insert(class.to_string());
                }
                let stripped = strip_locations(&err);
                messages.push(Message { role: "assistant".into(), content: program });
                messages.push(Message { role: "user".into(),      content: stripped });
            }
        }
    }

    let consumed = prompt_total + completion_total;
    Ok((
        ScoreRow {
            cohort: cohort.as_str().to_string(),
            task_id: task.id.clone(),
            first_attempt_green,
            turns_to_green,
            prompt_tokens: prompt_total,
            completion_tokens: completion_total,
            error_classes,
            final_status,
        },
        consumed,
    ))
}
