//! The `calibration:` check block.
//!
//! Each entry reads a labels file (one `{confidence, correct}` row per judged
//! item), computes Expected Calibration Error and the Brier score with the
//! existing [`crate::eval::jury::calibration`] math, and exposes them as the
//! assertable targets `ece` and `brier`. When the entry also supplies a
//! confusion matrix and an observed positive rate, it exposes the Rogan-Gladen
//! `corrected_rate` target so a user can check the bias-corrected pass rate is
//! no worse than the raw one, plus the `corrected_rate_low` /
//! `corrected_rate_high` targets, the Wald 95% interval around that
//! rate, so a user can gate on the band rather than the point estimate.
//!
//! The user produces the labels file by running their judge against a
//! hand-labeled set and recording, per item, the judge's stated confidence and
//! whether the verdict matched the trusted label. It is a pure file read plus
//! the existing math; no network and no model calls.

use std::path::{Path, PathBuf};

use serde::Deserialize;
use serde_json::{json, Value};
use thiserror::Error;

use crate::eval::jury::calibration::{
    brier_score, calibrated_calibration_error, calibration_error, corrected_rate_ci,
    estimate_reliability, fit_linear_probe, rogan_gladen, LabeledVerdict,
};
use crate::suite::Assertion;

/// One parsed `calibration:` entry.
///
/// The `expect` list is the user's assertions over the targets; an empty list
/// means "apply the sane defaults" (see [`default_expect`]). `reliability` and
/// `observed_positive_rate` are both required to enable the `corrected_rate`
/// target, so they are carried as a pair-or-nothing via two `Option`s the
/// evaluator reads together.
#[derive(Debug, Clone, PartialEq)]
pub struct CalibrationCheck {
    /// Human-readable name; the report row name for this entry.
    pub name: String,
    /// Path to the labels file (JSONL or a YAML array of `{confidence,
    /// correct}`), resolved relative to the suite file's directory.
    pub labels: PathBuf,
    /// Optional confusion matrix on the trusted set. Enables the
    /// `corrected_rate` target together with `observed_positive_rate`.
    pub reliability: Option<ConfusionMatrix>,
    /// Optional observed positive rate the judge reported on the full set.
    /// Enables the `corrected_rate` target together with `reliability`.
    pub observed_positive_rate: Option<f64>,
    /// Assertions over the targets. Empty means "apply the defaults".
    pub expect: Vec<Assertion>,
}

/// Confusion-matrix counts on a small trusted label set, the input to the
/// Rogan-Gladen bias correction.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ConfusionMatrix {
    /// True positives: truth positive, judge passed.
    pub tp: u32,
    /// False negatives: truth positive, judge failed.
    pub fn_: u32,
    /// True negatives: truth negative, judge failed.
    pub tn: u32,
    /// False positives: truth negative, judge passed.
    pub fp: u32,
}

/// One labeled verdict row in a labels file.
#[derive(Debug, Clone, Copy, Deserialize)]
struct LabelRow {
    confidence: f64,
    correct: bool,
}

/// Errors raised while loading a labels file.
///
/// A missing or unparseable labels file is a load error, the same class as a
/// malformed config: the suite cannot run a check whose input it cannot read.
#[derive(Debug, Error)]
pub enum CalibrationError {
    /// The labels file could not be read from disk.
    #[error("calibration `{name}`: could not read labels file `{path}`: {source}")]
    Read {
        /// Entry name for the error message.
        name: String,
        /// Path that could not be read.
        path: String,
        /// Underlying I/O error.
        source: std::io::Error,
    },
    /// The labels file was read but did not parse as JSONL or a YAML array of
    /// `{confidence, correct}` rows.
    #[error("calibration `{name}`: labels file `{path}` is not valid JSONL or a YAML array of {{confidence, correct}}: {detail}")]
    Parse {
        /// Entry name for the error message.
        name: String,
        /// Path that failed to parse.
        path: String,
        /// One-line parse-failure summary.
        detail: String,
    },
}

/// Default ECE ceiling when the entry declares no `expect:`. Matches the
/// design's sane default: stated confidence should track accuracy within ten
/// points.
const DEFAULT_ECE_MAX: f64 = 0.10;
/// Default Brier ceiling when the entry declares no `expect:`.
const DEFAULT_BRIER_MAX: f64 = 0.25;

/// Load a labels file as a list of [`LabeledVerdict`].
///
/// Accepts JSONL (one `{confidence, correct}` object per non-blank line) or a
/// YAML array of the same objects. A read failure or a parse failure is a load
/// error so a typo in the path or the file surfaces before the run, not as a
/// silent empty set.
pub fn load_labels(name: &str, path: &Path) -> Result<Vec<LabeledVerdict>, CalibrationError> {
    let raw = std::fs::read_to_string(path).map_err(|source| CalibrationError::Read {
        name: name.to_string(),
        path: path.display().to_string(),
        source,
    })?;
    let rows = parse_label_rows(&raw).map_err(|detail| CalibrationError::Parse {
        name: name.to_string(),
        path: path.display().to_string(),
        detail,
    })?;
    Ok(rows
        .into_iter()
        .map(|r| LabeledVerdict {
            confidence: r.confidence,
            correct: r.correct,
        })
        .collect())
}

/// Parse label rows from JSONL or a YAML array. Tries JSONL first (the
/// documented primary form): every non-blank line must be a `LabelRow`. If any
/// line fails as JSON, falls back to parsing the whole document as a YAML
/// array, so a `- {confidence: .., correct: ..}` list also loads.
fn parse_label_rows(raw: &str) -> Result<Vec<LabelRow>, String> {
    let non_blank: Vec<&str> = raw
        .lines()
        .map(str::trim)
        .filter(|l| !l.is_empty())
        .collect();
    let looks_like_jsonl = non_blank.iter().all(|l| l.starts_with('{'));
    if !non_blank.is_empty() && looks_like_jsonl {
        let mut rows = Vec::with_capacity(non_blank.len());
        for (i, line) in non_blank.iter().enumerate() {
            let row: LabelRow = serde_json::from_str(line)
                .map_err(|e| format!("line {} is not a valid row: {e}", i + 1))?;
            rows.push(row);
        }
        return Ok(rows);
    }
    // YAML array form (also parses a JSON array, since JSON is a YAML subset).
    serde_yaml::from_str::<Vec<LabelRow>>(raw).map_err(|e| e.to_string())
}

/// Build the targets object the `calibration:` assertions resolve against.
///
/// Shape: `{ "ece": <number>, "brier": <number>, "ece_calibrated": <number>,
/// "corrected_rate": <number|absent>, "corrected_rate_low": <number|absent>,
/// "corrected_rate_high": <number|absent> }`. `ece_calibrated` is the ECE after
/// a 1-D logistic probe is fit on the same labels, always present because it is
/// deterministic and additive. The three corrected-rate targets
/// are present only when both a confusion matrix and an observed positive rate
/// were supplied (the same gate); otherwise each resolves to null and an
/// assertion on it fails with a clear "absent" diff rather than a silent pass.
/// `corrected_rate_low` / `corrected_rate_high` are the Wald 95% interval
/// around `corrected_rate`.
pub fn metrics_object(check: &CalibrationCheck, samples: &[LabeledVerdict]) -> Value {
    let ece = calibration_error(samples);
    let brier = brier_score(samples);
    // Fit a 1-D logistic probe on the same labels and report the
    // post-calibration ECE, so a user can see whether a cheap recalibration
    // would let an over-confident judge gate honestly. Deterministic and
    // additive, so it is always present and needs no new YAML field.
    let ece_calibrated = calibrated_calibration_error(samples, &fit_linear_probe(samples));
    let mut obj = json!({ "ece": ece, "brier": brier, "ece_calibrated": ece_calibrated });
    if let Some(rate) = corrected_rate(check) {
        obj["corrected_rate"] = json!(rate);
    }
    if let Some((low, high)) = corrected_rate_band(check) {
        obj["corrected_rate_low"] = json!(low);
        obj["corrected_rate_high"] = json!(high);
    }
    obj
}

/// The Rogan-Gladen corrected positive rate, or `None` when the entry did not
/// supply both a confusion matrix and an observed positive rate.
fn corrected_rate(check: &CalibrationCheck) -> Option<f64> {
    let matrix = check.reliability?;
    let observed = check.observed_positive_rate?;
    let reliability = estimate_reliability(matrix.tp, matrix.fn_, matrix.tn, matrix.fp);
    Some(rogan_gladen(observed, &reliability))
}

/// The Wald 95% confidence interval `(low, high)` around the corrected rate
///, or `None` under the same gate as [`corrected_rate`].
///
/// The trusted-set size is the total of the confusion-matrix counts, which is
/// the binomial sample size the interval widens with: a larger hand-labeled
/// set yields a tighter band.
fn corrected_rate_band(check: &CalibrationCheck) -> Option<(f64, f64)> {
    let matrix = check.reliability?;
    let observed = check.observed_positive_rate?;
    let reliability = estimate_reliability(matrix.tp, matrix.fn_, matrix.tn, matrix.fp);
    let trusted_set_size = matrix.tp + matrix.fn_ + matrix.tn + matrix.fp;
    Some(corrected_rate_ci(observed, &reliability, trusted_set_size))
}

/// The default assertion set applied when an entry declares no `expect:`.
///
/// `ece <= 0.10` and `brier <= 0.25` always apply. When the corrected-rate
/// inputs are present, a third default asserts `corrected_rate <=
/// observed_positive_rate`, the "no worse than the raw rate" check the design
/// calls for. Returns the assertions as the same [`Assertion`] type a user
/// would write, so the evaluation path is identical for default and explicit
/// `expect:`.
pub fn default_expect(check: &CalibrationCheck) -> Vec<Assertion> {
    let mut out = vec![
        metric_assertion("ece", json!({ "maximum": DEFAULT_ECE_MAX })),
        metric_assertion("brier", json!({ "maximum": DEFAULT_BRIER_MAX })),
    ];
    if let Some(observed) = check.observed_positive_rate {
        if check.reliability.is_some() {
            out.push(metric_assertion(
                "corrected_rate",
                json!({ "maximum": observed }),
            ));
        }
    }
    out
}

/// Build a `schema`-matcher assertion over a numeric target. Reuses the
/// existing assertion grammar so a default and a user-written `expect:` are
/// the same shape downstream.
fn metric_assertion(target: &str, schema: Value) -> Assertion {
    use crate::suite::Matcher;
    Assertion {
        target: target.to_string(),
        matcher: Matcher::Schema(schema),
        weight: 1.0,
        name: None,
        transform: None,
    }
}

/// Evaluate one calibration entry end to end.
///
/// Loads the labels, builds the targets, applies the entry's `expect:` (or the
/// defaults when it declared none), and folds the result into a
/// [`crate::checks::CheckOutcome`]. An empty labels set is not an error (ECE is
/// 0 by definition) but attaches a warning note so it does not read as a silent
/// pass. A label-load failure is returned as [`CalibrationError`] so the caller
/// can fail the suite load.
pub fn evaluate(
    check: &CalibrationCheck,
    base_dir: &Path,
) -> Result<super::CheckOutcome, CalibrationError> {
    let path = resolve_path(base_dir, &check.labels);
    let samples = load_labels(&check.name, &path)?;
    let metrics = metrics_object(check, &samples);

    let assertions: Vec<Assertion> = if check.expect.is_empty() {
        default_expect(check)
    } else {
        check.expect.clone()
    };

    let mut failure: Option<String> = None;
    for (i, assertion) in assertions.iter().enumerate() {
        if let Err(reason) = crate::executor::eval_metric_assertion(&metrics, assertion) {
            if failure.is_none() {
                failure = Some(format!(
                    "assertion #{i} (`{}`) failed: {reason}",
                    assertion.target
                ));
            }
        }
    }

    // An empty labels set passes (ECE is 0) but should never read as a silent
    // pass, so attach the warning even on a clean run.
    if samples.is_empty() && failure.is_none() {
        return Ok(super::CheckOutcome {
            name: check.name.clone(),
            passed: true,
            failure: Some(
                "warning: labels file is empty; ECE and Brier are 0 by definition".to_string(),
            ),
        });
    }

    Ok(super::CheckOutcome {
        name: check.name.clone(),
        passed: failure.is_none(),
        failure,
    })
}

/// Resolve a labels path against the suite file's directory unless it is
/// already absolute, so relative paths in the YAML are read relative to the
/// suite rather than the process cwd.
fn resolve_path(base_dir: &Path, path: &Path) -> PathBuf {
    if path.is_absolute() {
        path.to_path_buf()
    } else {
        base_dir.join(path)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;
    use tempfile::NamedTempFile;

    fn check(expect: Vec<Assertion>) -> CalibrationCheck {
        CalibrationCheck {
            name: "judge".into(),
            labels: PathBuf::from("labels.jsonl"),
            reliability: None,
            observed_positive_rate: None,
            expect,
        }
    }

    #[test]
    fn parses_jsonl_rows() {
        let raw =
            "{\"confidence\": 0.9, \"correct\": true}\n{\"confidence\": 0.2, \"correct\": false}\n";
        let rows = parse_label_rows(raw).expect("jsonl parses");
        assert_eq!(rows.len(), 2);
        assert!(rows[0].correct);
        assert!(!rows[1].correct);
    }

    #[test]
    fn parses_yaml_array() {
        let raw = "- confidence: 0.9\n  correct: true\n- confidence: 0.1\n  correct: false\n";
        let rows = parse_label_rows(raw).expect("yaml array parses");
        assert_eq!(rows.len(), 2);
    }

    #[test]
    fn malformed_jsonl_line_is_an_error() {
        let raw = "{\"confidence\": 0.9, \"correct\": true}\n{not json}\n";
        assert!(parse_label_rows(raw).is_err());
    }

    #[test]
    fn missing_labels_file_is_a_load_error() {
        let err = load_labels("judge", Path::new("/no/such/labels.jsonl"))
            .expect_err("missing file errors");
        assert!(matches!(err, CalibrationError::Read { .. }));
    }

    #[test]
    fn default_expect_has_ece_and_brier_without_corrected_rate() {
        let c = check(Vec::new());
        let defaults = default_expect(&c);
        let targets: Vec<&str> = defaults.iter().map(|a| a.target.as_str()).collect();
        assert_eq!(targets, vec!["ece", "brier"]);
    }

    #[test]
    fn default_expect_adds_corrected_rate_when_inputs_present() {
        let mut c = check(Vec::new());
        c.reliability = Some(ConfusionMatrix {
            tp: 90,
            fn_: 10,
            tn: 80,
            fp: 20,
        });
        c.observed_positive_rate = Some(0.5);
        let defaults = default_expect(&c);
        let targets: Vec<&str> = defaults.iter().map(|a| a.target.as_str()).collect();
        assert_eq!(targets, vec!["ece", "brier", "corrected_rate"]);
    }

    #[test]
    fn corrected_rate_target_absent_without_inputs() {
        let c = check(Vec::new());
        let obj = metrics_object(&c, &[]);
        assert!(obj.get("corrected_rate").is_none());
    }

    #[test]
    fn metrics_object_emits_calibrated_ece_no_worse_than_raw_for_overconfident_judge() {
        // An over-confident judge: states 0.95 but is correct ~60% of the
        // time. The calibrated ECE is always present and, after fitting the
        // probe, should be no worse than the raw ECE.
        let mut samples = Vec::with_capacity(100);
        for i in 0..100 {
            samples.push(LabeledVerdict {
                confidence: 0.95,
                correct: i < 60,
            });
        }
        let c = check(Vec::new());
        let obj = metrics_object(&c, &samples);
        let ece = obj["ece"].as_f64().expect("ece");
        let ece_calibrated = obj["ece_calibrated"].as_f64().expect("ece_calibrated");
        assert!(
            ece_calibrated <= ece,
            "calibrated ECE {ece_calibrated} should be <= raw ECE {ece}"
        );
    }

    #[test]
    fn corrected_rate_ci_targets_absent_without_inputs() {
        // The band shares corrected_rate's gate: no matrix, no band.
        let c = check(Vec::new());
        let obj = metrics_object(&c, &[]);
        assert!(obj.get("corrected_rate_low").is_none());
        assert!(obj.get("corrected_rate_high").is_none());
    }

    #[test]
    fn corrected_rate_ci_targets_present_and_bracket_the_point() {
        // With the matrix and observed rate, the band targets resolve and
        // bracket the point estimate.
        let mut c = check(Vec::new());
        c.reliability = Some(ConfusionMatrix {
            tp: 90,
            fn_: 10,
            tn: 80,
            fp: 20,
        });
        c.observed_positive_rate = Some(0.5);
        let obj = metrics_object(&c, &[]);
        let point = obj["corrected_rate"].as_f64().expect("corrected_rate");
        let low = obj["corrected_rate_low"].as_f64().expect("low");
        let high = obj["corrected_rate_high"].as_f64().expect("high");
        assert!(
            low <= point && point <= high,
            "low {low} <= {point} <= high {high}"
        );
        assert!(
            low < high,
            "a non-empty trusted set yields a positive-width band"
        );
    }

    fn write_labels(body: &str) -> NamedTempFile {
        let mut f = NamedTempFile::new().unwrap();
        f.write_all(body.as_bytes()).unwrap();
        f
    }

    #[test]
    fn well_calibrated_judge_passes_the_defaults() {
        // Confident-and-right plus unconfident-and-wrong: ECE and Brier are 0.
        let f = write_labels(
            "{\"confidence\": 1.0, \"correct\": true}\n{\"confidence\": 0.0, \"correct\": false}\n",
        );
        let mut c = check(Vec::new());
        c.labels = f.path().to_path_buf();
        let outcome = evaluate(&c, Path::new("/")).expect("evaluates");
        assert!(outcome.passed, "{outcome:?}");
    }

    #[test]
    fn overconfident_judge_fails_the_ece_default() {
        // Always claims 1.0 confidence but right only half the time: ECE 0.5,
        // well over the 0.10 default ceiling.
        let f = write_labels(
            "{\"confidence\": 1.0, \"correct\": true}\n{\"confidence\": 1.0, \"correct\": false}\n",
        );
        let mut c = check(Vec::new());
        c.labels = f.path().to_path_buf();
        let outcome = evaluate(&c, Path::new("/")).expect("evaluates");
        assert!(!outcome.passed);
        assert!(outcome.failure.unwrap().contains("ece"));
    }

    #[test]
    fn empty_labels_set_passes_with_a_warning() {
        let f = write_labels("");
        let mut c = check(Vec::new());
        c.labels = f.path().to_path_buf();
        let outcome = evaluate(&c, Path::new("/")).expect("evaluates");
        assert!(outcome.passed);
        assert!(outcome.failure.unwrap().contains("empty"));
    }

    #[test]
    fn corrected_rate_default_holds_when_correction_helps() {
        // A judge that over-reports positives: observed 0.5, corrected lands
        // below it, so corrected_rate <= observed_positive_rate passes. The
        // labels are perfectly calibrated so ECE/Brier are 0 and only the
        // corrected-rate default is under test.
        let f = write_labels(
            "{\"confidence\": 1.0, \"correct\": true}\n{\"confidence\": 0.0, \"correct\": false}\n",
        );
        let mut c = check(Vec::new());
        c.labels = f.path().to_path_buf();
        c.reliability = Some(ConfusionMatrix {
            tp: 90,
            fn_: 10,
            tn: 80,
            fp: 20,
        });
        c.observed_positive_rate = Some(0.5);
        let outcome = evaluate(&c, Path::new("/")).expect("evaluates");
        // The corrected rate (~0.43) is below the observed 0.5.
        assert!(outcome.passed, "{outcome:?}");
    }
}
