//! Measure a decision backend against labelled cases.
//!
//! The judge got an eval before it got trusted (`scripts/judge-eval.sh`);
//! a triage model gets the same. Every backend answers the SAME Noul per
//! (task, skill) pair, and is scored on ranking (AUROC), on the operating
//! point (best-F1 threshold and F1 at 0.5), and on calibration (Brier, ECE)
//! — because a calibrated middle band is the whole reason to have this tier,
//! and a backend that ranks well but says 0.9 to everything has none.
//!
//!   decide-eval [--kind skills|door] [--backend jev|nli|lexical|all]
//!               [--wording named|plain] [--set path] [--skills dir] [--dump dir]
//!
//! `--dump` writes every pair's probability so a temperature can be fitted
//! offline without re-running the model.

use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::time::Instant;

use cm_decide::{Answer, DecideError, Decider, Decision, Question};
use serde::{Deserialize, Serialize};

#[derive(Deserialize)]
struct EvalSet {
    cases: Vec<Case>,
}

#[derive(Deserialize, Clone)]
struct Case {
    id: String,
    task: String,
    positives: Vec<String>,
}

#[derive(Deserialize)]
struct DoorSet {
    cases: Vec<DoorCase>,
}

#[derive(Deserialize, Clone)]
struct DoorCase {
    id: String,
    tool: String,
    args: serde_json::Value,
    deny: bool,
}

/// The door governor: one call per action, three Nouls, the max is the deny
/// probability. Scored like the skill pairs (label = deny), plus the three-
/// band outcome at the shipped thresholds — because what ships is not a
/// threshold but a band, and the number an operator needs is how many
/// actions the model would decide alone and how many it would hand over.
async fn run_door(backend: &dyn Decider, cases: &[DoorCase]) -> (Report, [usize; 6]) {
    use cm_decide::door::{ALLOW_BELOW, DENY_AT};
    let mut r = Report::default();
    // [allow-labelled → allowed, → review, → denied, deny-labelled → allowed, → review, → denied]
    let mut bands = [0usize; 6];
    let questions = cm_decide::door::questions();
    for case in cases {
        let state = cm_decide::door::state(&case.tool, &case.args);
        let d = match backend.decide(&state, &questions).await {
            Ok(d) => d,
            Err(e) => {
                eprintln!("  {}: {} failed: {e}", backend.name(), case.id);
                r.errors += 1;
                continue;
            }
        };
        r.latency_ms.push(d.latency.as_millis());
        if let Some(u) = d.usage {
            r.input_tokens += u.input_tokens;
        }
        let Some(risk) = cm_decide::door::Risk::from_answers(&d.answers) else {
            r.errors += 1;
            continue;
        };
        r.pairs.push(Pair { case: case.id.clone(), skill: "deny".into(), label: case.deny, p: risk.deny });
        let band = if risk.deny >= DENY_AT { 2 } else if risk.deny < ALLOW_BELOW { 0 } else { 1 };
        bands[if case.deny { 3 } else { 0 } + band] += 1;
        println!(
            "  {:<32} {} deny {:.2}  (exfil {:.2} secret {:.2} spam {:.2}) → {}",
            case.id,
            if case.deny { "DENY " } else { "allow" },
            risk.deny, risk.exfil, risk.secret, risk.spam,
            ["allow", "REVIEW", "deny"][band]
        );
    }
    (r, bands)
}

#[derive(Clone)]
struct Skill {
    name: String,
    when_to_use: String,
}

/// The question a backend gets. `--wording` picks the template; a backend
/// is reported under the wording it was run with, and each backend ships
/// with the wording that measured best FOR IT — a prompt is part of the
/// backend, and an NLI cross-encoder and a decision model do not want the
/// same sentence. Both are run on both so the table says so.
fn question_for(skill: &Skill, wording: &str) -> Question {
    let text = match wording {
        // The shipped question, shared with the server's shadow path.
        "named" => return cm_decide::triage::question(&skill.name, &skill.when_to_use),
        // The when_to_use line alone, second person rewritten to the agent,
        // as a plain declarative the NLI head was trained on.
        "plain" => third_person(&skill.when_to_use),
        other => panic!("unknown wording {other}"),
    };
    Question::noul(text)
}

/// "You're the tester on a mobile team" → "The agent is the tester on a
/// mobile team". Crude on purpose: it is a rewrite of a dozen fixed
/// openings, not a grammar, and it exists so the NLI hypothesis reads like
/// an MNLI hypothesis.
fn third_person(when: &str) -> String {
    let w = when.trim();
    let rules: &[(&str, &str)] = &[
        ("You're the ", "The agent is the "),
        ("You are the ", "The agent is the "),
        ("You're a ", "The agent is a "),
        ("You are a ", "The agent is a "),
        ("You're on a ", "The agent is on a "),
        ("You are on a ", "The agent is on a "),
        ("You're ", "The agent is "),
        ("You are ", "The agent is "),
        ("You need ", "The agent needs "),
        ("You have ", "The agent has "),
    ];
    for (from, to) in rules {
        if let Some(rest) = w.strip_prefix(from) {
            return format!("{to}{rest}");
        }
    }
    format!("In this task, {}", w.trim_end_matches('.').to_string() + ".")
}

fn load_skills(dir: &Path) -> Vec<Skill> {
    let mut out = Vec::new();
    fn walk(dir: &Path, out: &mut Vec<Skill>) {
        let Ok(rd) = std::fs::read_dir(dir) else { return };
        let mut entries: Vec<_> = rd.flatten().collect();
        entries.sort_by_key(|e| e.path());
        for e in entries {
            let p = e.path();
            if p.is_dir() {
                walk(&p, out);
            } else if p.extension().is_some_and(|x| x == "md") {
                let Ok(text) = std::fs::read_to_string(&p) else { continue };
                let field = |k: &str| {
                    text.lines()
                        .find_map(|l| l.strip_prefix(k).map(|v| v.trim().trim_matches('"').to_string()))
                };
                if let (Some(name), Some(when)) = (field("name:"), field("when_to_use:")) {
                    out.push(Skill { name, when_to_use: when });
                }
            }
        }
    }
    walk(dir, &mut out);
    out
}

/// Keyword overlap between the task and the skill's name + when_to_use:
/// the bar a model has to clear to be worth a network call.
struct Lexical;

fn tokens(s: &str) -> std::collections::BTreeSet<String> {
    s.to_ascii_lowercase()
        .split(|c: char| !c.is_alphanumeric())
        .filter(|w| w.len() > 3)
        .map(str::to_string)
        .collect()
}

#[async_trait::async_trait]
impl Decider for Lexical {
    fn name(&self) -> &str {
        "lexical"
    }
    async fn decide(
        &self,
        state: &str,
        questions: &BTreeMap<String, Question>,
    ) -> Result<Decision, DecideError> {
        let started = Instant::now();
        let st = tokens(state);
        let answers = questions
            .iter()
            .map(|(id, q)| {
                let Question::Noul { instructions, .. } = q else { unreachable!() };
                let ht = tokens(instructions);
                let inter = st.intersection(&ht).count() as f64;
                let p = if ht.is_empty() { 0.0 } else { (inter / ht.len() as f64 * 4.0).min(1.0) };
                (id.clone(), Answer::Noul { noul: p })
            })
            .collect();
        Ok(Decision { model: "lexical".into(), answers, usage: None, latency: started.elapsed() })
    }
}

#[derive(Serialize)]
struct Pair {
    case: String,
    skill: String,
    label: bool,
    p: f64,
}

#[derive(Default)]
struct Report {
    pairs: Vec<Pair>,
    latency_ms: Vec<u128>,
    input_tokens: u64,
    per_case_topk_hits: Vec<(usize, usize)>,
    errors: usize,
}

fn auroc(pairs: &[Pair]) -> f64 {
    // Mann–Whitney: fraction of (positive, negative) pairs ranked correctly.
    let pos: Vec<f64> = pairs.iter().filter(|p| p.label).map(|p| p.p).collect();
    let neg: Vec<f64> = pairs.iter().filter(|p| !p.label).map(|p| p.p).collect();
    if pos.is_empty() || neg.is_empty() {
        return f64::NAN;
    }
    let mut s = 0.0;
    for a in &pos {
        for b in &neg {
            s += if a > b { 1.0 } else if a == b { 0.5 } else { 0.0 };
        }
    }
    s / (pos.len() * neg.len()) as f64
}

fn f1_at(pairs: &[Pair], t: f64) -> (f64, f64, f64) {
    let (mut tp, mut fp, mut fn_) = (0.0, 0.0, 0.0);
    for p in pairs {
        match (p.p >= t, p.label) {
            (true, true) => tp += 1.0,
            (true, false) => fp += 1.0,
            (false, true) => fn_ += 1.0,
            _ => {}
        }
    }
    let prec = if tp + fp > 0.0 { tp / (tp + fp) } else { 0.0 };
    let rec = if tp + fn_ > 0.0 { tp / (tp + fn_) } else { 0.0 };
    let f1 = if prec + rec > 0.0 { 2.0 * prec * rec / (prec + rec) } else { 0.0 };
    (f1, prec, rec)
}

fn brier(pairs: &[Pair]) -> f64 {
    pairs.iter().map(|p| (p.p - if p.label { 1.0 } else { 0.0 }).powi(2)).sum::<f64>() / pairs.len() as f64
}

/// Expected calibration error, ten equal-width bins.
fn ece(pairs: &[Pair]) -> f64 {
    let mut bins = vec![(0usize, 0.0f64, 0.0f64); 10];
    for p in pairs {
        let b = ((p.p * 10.0).floor() as usize).min(9);
        bins[b].0 += 1;
        bins[b].1 += p.p;
        bins[b].2 += if p.label { 1.0 } else { 0.0 };
    }
    let n = pairs.len() as f64;
    bins.iter()
        .filter(|(c, _, _)| *c > 0)
        .map(|(c, sp, sy)| (*c as f64 / n) * ((sp / *c as f64) - (sy / *c as f64)).abs())
        .sum()
}

async fn run(backend: &dyn Decider, cases: &[Case], skills: &[Skill], wording: &str) -> Report {
    let mut r = Report::default();
    let questions: BTreeMap<String, Question> =
        skills.iter().map(|s| (s.name.clone(), question_for(s, wording))).collect();
    for case in cases {
        let d = match backend.decide(&case.task, &questions).await {
            Ok(d) => d,
            Err(e) => {
                eprintln!("  {}: {} failed: {e}", backend.name(), case.id);
                r.errors += 1;
                continue;
            }
        };
        r.latency_ms.push(d.latency.as_millis());
        if let Some(u) = d.usage {
            r.input_tokens += u.input_tokens;
        }
        let mut ranked: Vec<(String, f64)> = Vec::new();
        for s in skills {
            let p = match d.answers.get(&s.name) {
                Some(Answer::Noul { noul }) => *noul,
                _ => 0.0,
            };
            let label = case.positives.iter().any(|x| x == &s.name);
            r.pairs.push(Pair { case: case.id.clone(), skill: s.name.clone(), label, p });
            ranked.push((s.name.clone(), p));
        }
        ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
        let k = case.positives.len();
        let hits = ranked.iter().take(k).filter(|(n, _)| case.positives.contains(n)).count();
        r.per_case_topk_hits.push((hits, k));
    }
    r
}

fn print_report(name: &str, r: &Report) {
    if r.pairs.is_empty() {
        println!("{name:<24} no results ({} errors)", r.errors);
        return;
    }
    let best = (5..=95)
        .map(|i| i as f64 / 100.0)
        .map(|t| (t, f1_at(&r.pairs, t)))
        .max_by(|a, b| a.1 .0.partial_cmp(&b.1 .0).unwrap())
        .unwrap();
    let (f1_half, p_half, r_half) = f1_at(&r.pairs, 0.5);
    let topk: (usize, usize) = r.per_case_topk_hits.iter().fold((0, 0), |a, b| (a.0 + b.0, a.1 + b.1));
    let lat = if r.latency_ms.is_empty() { 0 } else { r.latency_ms.iter().sum::<u128>() / r.latency_ms.len() as u128 };
    println!(
        "{name:<24} AUROC {:.3}  F1@0.5 {:.2} (P {:.2} R {:.2})  bestF1 {:.2}@{:.2}  top-k {}/{}  Brier {:.3}  ECE {:.3}  {} ms/call  {} tok  errors {}",
        auroc(&r.pairs), f1_half, p_half, r_half, best.1 .0, best.0, topk.0, topk.1,
        brier(&r.pairs), ece(&r.pairs), lat, r.input_tokens, r.errors
    );
}

#[tokio::main]
async fn main() {
    let mut args = std::env::args().skip(1);
    let (mut backend, mut set, mut skills_dir, mut dump, mut wording, mut kind) = (
        "all".to_string(),
        PathBuf::from("crates/cm-decide/eval/skill-triage.json"),
        PathBuf::from("skills"),
        None::<PathBuf>,
        "named".to_string(),
        "skills".to_string(),
    );
    while let Some(a) = args.next() {
        match a.as_str() {
            "--backend" => backend = args.next().unwrap_or_default(),
            "--kind" => {
                kind = args.next().unwrap_or_default();
                if kind == "door" && set.ends_with("skill-triage.json") {
                    set = PathBuf::from("crates/cm-decide/eval/door-actions.json");
                }
            }
            "--wording" => wording = args.next().unwrap_or_default(),
            "--set" => set = args.next().map(PathBuf::from).unwrap_or(set),
            "--skills" => skills_dir = args.next().map(PathBuf::from).unwrap_or(skills_dir),
            "--dump" => dump = args.next().map(PathBuf::from),
            other => {
                eprintln!("unknown arg {other}");
                std::process::exit(2);
            }
        }
    }
    let text = std::fs::read_to_string(&set).expect("read eval set");
    if kind == "door" {
        let door: DoorSet = serde_json::from_str(&text).expect("parse door set");
        println!(
            "{} door actions, {} labelled deny; bands: allow < {} ≤ review < {} ≤ deny\n",
            door.cases.len(),
            door.cases.iter().filter(|c| c.deny).count(),
            cm_decide::door::ALLOW_BELOW,
            cm_decide::door::DENY_AT
        );
        let mut backends: Vec<Box<dyn Decider>> = Vec::new();
        #[cfg(feature = "jev")]
        if backend == "all" || backend == "jev" {
            match cm_decide::jev::Jev::from_env() {
                Some(j) => backends.push(Box::new(j)),
                None => eprintln!("jev: TYPESAFE_API_KEY unset — skipped"),
            }
        }
        #[cfg(feature = "nli")]
        if backend == "all" || backend == "nli" {
            match cm_decide::nli::Nli::load(cm_decide::nli::NliConfig::from_env()) {
                Ok(n) => backends.push(Box::new(n)),
                Err(e) => eprintln!("nli: could not load — {e}"),
            }
        }
        for b in &backends {
            let (r, bands) = run_door(b.as_ref(), &door.cases).await;
            println!();
            print_report(b.name(), &r);
            println!(
                "{:<24} allow-labelled: {} allowed / {} review / {} DENIED (false denies)   deny-labelled: {} ALLOWED (misses) / {} review / {} denied",
                "", bands[0], bands[1], bands[2], bands[3], bands[4], bands[5]
            );
        }
        return;
    }
    let set: EvalSet = serde_json::from_str(&text).expect("parse eval set");
    let skills = load_skills(&skills_dir);
    // Every positive must name a real skill, or the label is a typo scored
    // as a miss against every backend.
    for c in &set.cases {
        for p in &c.positives {
            assert!(skills.iter().any(|s| &s.name == p), "case {}: unknown skill {p:?}", c.id);
        }
    }
    println!(
        "{} cases × {} skills = {} pairs, {} positive\n",
        set.cases.len(),
        skills.len(),
        set.cases.len() * skills.len(),
        set.cases.iter().map(|c| c.positives.len()).sum::<usize>()
    );

    let mut backends: Vec<Box<dyn Decider>> = Vec::new();
    if backend == "all" || backend == "lexical" {
        backends.push(Box::new(Lexical));
    }
    #[cfg(feature = "jev")]
    if backend == "all" || backend == "jev" {
        match cm_decide::jev::Jev::from_env() {
            Some(j) => backends.push(Box::new(j)),
            None => eprintln!("jev: TYPESAFE_API_KEY unset — skipped"),
        }
    }
    #[cfg(feature = "nli")]
    if backend == "all" || backend == "nli" {
        let t0 = Instant::now();
        match cm_decide::nli::Nli::load(cm_decide::nli::NliConfig::from_env()) {
            Ok(n) => {
                eprintln!("nli: loaded {} in {:?}", n.name(), t0.elapsed());
                backends.push(Box::new(n));
            }
            Err(e) => eprintln!("nli: could not load — {e}"),
        }
    }
    if backends.is_empty() {
        eprintln!("no backend to run");
        std::process::exit(2);
    }
    for b in &backends {
        let r = run(b.as_ref(), &set.cases, &skills, &wording).await;
        print_report(&format!("{} [{wording}]", b.name()), &r);
        if let Some(dir) = &dump {
            std::fs::create_dir_all(dir).ok();
            let path = dir.join(format!("{}-{wording}.json", b.name().replace([':', '/'], "-")));
            std::fs::write(&path, serde_json::to_string_pretty(&r.pairs).unwrap()).ok();
        }
    }
}