//! Bias-mitigation wiring for the `llm-jury` matcher (WOR-1296).
//!
//! The pure mitigation math lives in [`crate::eval::jury::bias`]; this
//! module is the matcher-side glue that turns a parsed `bias:` YAML block
//! into a prompt edit, an anchor calibration pre-pass, and jury-note
//! annotations. It sits next to [`super::llm_jury`] instead of inside it so
//! that module stays within the size guard.
//!
//! Three of the five published mitigations are wired here: chain-of-thought
//! prompting, the verbosity guard, and rubric calibration against an anchor
//! set. Position swap's pairwise dispatch landed with the jury engine
//! convergence (WOR-1308, `crate::eval::jury::engine::DispatchMode`); what
//! still blocks a `position_swap:` knob is the matcher contract, which
//! carries one candidate where pairwise grading needs two. That YAML wiring
//! is the follow-up noted on WOR-1296. Same-family detection was wired
//! earlier (WOR-1299) and needs no YAML knob.

use serde::Deserialize;
use serde_json::Value;

use super::llm_judge::{self, JudgeConfig};
use super::llm_jury::{extract_score, JuryConfig, JuryJuror};
use crate::eval::jury::bias::{self, Anchor, CalibrationReport, CalibrationResult};
use crate::eval::provider::LlmProvider;
use crate::Error;

/// Parsed `bias:` block of an `llm-jury:` matcher.
///
/// Every mitigation defaults to off so configs written before WOR-1296 keep
/// byte-identical prompts and notes; users opt in per mitigation. Anchors
/// are real provider calls (one per juror per anchor, before the candidate
/// is graded), so an anchor set multiplies the jury's paid fan-out and the
/// cost must be a deliberate choice.
#[derive(Debug, Clone, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct BiasConfig {
    /// When true, juror prompts ask for step-by-step reasoning before the
    /// verdict (Wei et al., 2022). Off by default because reasoning costs
    /// output tokens on every juror call.
    #[serde(default)]
    pub chain_of_thought: bool,
    /// Token-ratio threshold for the verbosity guard (Saito et al., 2023).
    /// When the longer of the candidate and a juror's reply is at least
    /// this many times the shorter, the jury note carries a warning. Must
    /// be greater than 1.0; the published recommendation is 2.0.
    #[serde(default)]
    pub verbosity_guard: Option<f64>,
    /// Formatting-density threshold for the style-bias guard, in markers per
    /// 100 characters (arXiv:2604.23178). When the candidate's markdown and
    /// formatting density reaches this, the jury note carries a warning. Style
    /// bias now outweighs position bias as a judge confound, and it is a
    /// separate axis from verbosity. Must be greater than 0; around 5 is dense.
    #[serde(default)]
    pub style_guard: Option<f64>,
    /// Known-answer anchor set for rubric calibration (Kim et al., 2024).
    /// Jurors whose anchor score deviates from the expected score by more
    /// than [`ANCHOR_DEVIATION_THRESHOLD`] are flagged in the jury note.
    #[serde(default)]
    pub anchors: Vec<BiasAnchor>,
    /// When true, each juror also scores a formatting-heavy rendering of the
    /// candidate and the verdict drift is reported as `jury.style_drift`
    /// (WOR-1370, arXiv:2604.23178). Off by default because it doubles the
    /// per-juror provider calls; style bias now outweighs position bias, so
    /// the probe measures how much a verdict moves on formatting alone.
    #[serde(default)]
    pub style_swap: bool,
}

/// One known-answer anchor: a candidate plus the score a calibrated juror
/// should give it.
#[derive(Debug, Clone, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct BiasAnchor {
    /// Anchor candidate, graded through the same path as the real
    /// candidate. A string or any JSON value.
    pub candidate: Value,
    /// Score a calibrated juror should assign, on the 0 to 1 scale.
    pub expected: f64,
}

/// Maximum absolute deviation a juror may show on any anchor before it is
/// flagged as miscalibrated. Matches the construction-site default that
/// [`bias::RubricCalibration`] documents; not a YAML knob yet because no
/// ticket has asked for per-suite tuning.
pub const ANCHOR_DEVIATION_THRESHOLD: f64 = 0.2;

/// Instruction appended to the rubric when `chain_of_thought` is on.
///
/// Same wording the jury engine's prompt builder uses, so the two CoT paths
/// read identically to a juror. Riding on the rubric (instead of the judge
/// prompt template) keeps `PROMPT_TEMPLATE_VERSION` semantics intact: the
/// template text never changes, so cached verdicts stay valid.
pub(crate) const CHAIN_OF_THOUGHT_INSTRUCTION: &str =
    "Think step by step before deciding, but only return the JSON object in your reply.";

/// Reject out-of-range bias knobs at parse time, before any provider call.
///
/// Split from serde so the messages can name the offending value the same
/// way the sibling jury validations do.
pub(crate) fn validate(bias: &BiasConfig) -> crate::Result<()> {
    if let Some(ratio) = bias.verbosity_guard {
        if !ratio.is_finite() || ratio <= 1.0 {
            return Err(Error::InvalidMatcherInput(format!(
                "llm-jury bias verbosity_guard {ratio} must be a finite ratio greater than 1.0"
            )));
        }
    }
    if let Some(density) = bias.style_guard {
        if !density.is_finite() || density <= 0.0 {
            return Err(Error::InvalidMatcherInput(format!(
                "llm-jury bias style_guard {density} must be a finite density greater than 0"
            )));
        }
    }
    for anchor in &bias.anchors {
        if !anchor.expected.is_finite() || !(0.0..=1.0).contains(&anchor.expected) {
            return Err(Error::InvalidMatcherInput(format!(
                "llm-jury bias anchor expected score {} must be in [0, 1]",
                anchor.expected
            )));
        }
    }
    Ok(())
}

/// The rubric every juror actually sees: the configured rubric, plus the
/// chain-of-thought instruction when that mitigation is on.
///
/// Why the rubric and not the prompt template: appending here leaves the
/// `llm-judge` template (and its version contract) untouched, so a jury
/// with CoT off produces byte-identical prompts to before WOR-1296.
pub(crate) fn effective_rubric(rubric: &str, bias: Option<&BiasConfig>) -> String {
    match bias {
        Some(b) if b.chain_of_thought => format!("{rubric}\n\n{CHAIN_OF_THOUGHT_INSTRUCTION}"),
        _ => rubric.to_string(),
    }
}

/// Short tags for the mitigations a config turns on, rendered into the jury
/// summary note as `bias=[...]` so a reader can tell which protections were
/// active on a given verdict.
pub(crate) fn active_tags(bias: Option<&BiasConfig>) -> Vec<&'static str> {
    let Some(b) = bias else {
        return Vec::new();
    };
    let mut tags = Vec::new();
    if b.chain_of_thought {
        tags.push("cot");
    }
    if b.verbosity_guard.is_some() {
        tags.push("verbosity_guard");
    }
    if b.style_guard.is_some() {
        tags.push("style_guard");
    }
    if b.style_swap {
        tags.push("style_swap");
    }
    if !b.anchors.is_empty() {
        tags.push("anchors");
    }
    tags
}

/// Per-run bias outcome the jury fold renders into the summary note.
///
/// Carries which mitigations were configured plus the signals they raised,
/// so the note can show both "what was on" and "what fired" without the
/// fold re-deriving either.
pub(crate) struct BiasReport {
    /// Tags from [`active_tags`]; empty means no mitigation was on.
    pub(crate) active: Vec<&'static str>,
    /// Warning from the verbosity guard, when it tripped.
    pub(crate) verbosity_warning: Option<String>,
    /// Warning from the style-bias guard, when it tripped.
    pub(crate) style_warning: Option<String>,
    /// Anchor calibration result, when an anchor set was configured.
    pub(crate) calibration: Option<CalibrationReport>,
}

impl BiasReport {
    /// Suffix appended to the jury summary note, or `None` when no
    /// mitigation was configured (keeping pre-WOR-1296 notes byte-stable).
    pub(crate) fn render(&self) -> Option<String> {
        if self.active.is_empty() {
            return None;
        }
        let mut out = format!(" bias=[{}]", self.active.join(","));
        if let Some(warning) = &self.verbosity_warning {
            out.push_str(&format!("; verbosity: {warning}"));
        }
        if let Some(warning) = &self.style_warning {
            out.push_str(&format!("; style: {warning}"));
        }
        if let Some(report) = &self.calibration {
            if !report.miscalibrated_jurors.is_empty() {
                out.push_str(&format!(
                    "; miscalibrated jurors: {} (anchor deviation > {ANCHOR_DEVIATION_THRESHOLD})",
                    report.miscalibrated_jurors.join(", ")
                ));
            }
        }
        Some(out)
    }
}

/// Run the candidate-vs-reply verbosity guard over every juror reply.
///
/// `replies` pairs each juror's model with its reply note. Returns the
/// first warning [`bias::verbosity_warning`] raises, prefixed with the
/// juror it fired on; one warning is enough for the note, and the per-juror
/// rows already show every reply.
pub(crate) fn verbosity_signal(
    bias: Option<&BiasConfig>,
    candidate: &Value,
    replies: &[(&str, Option<&str>)],
) -> Option<String> {
    let threshold = bias?.verbosity_guard?;
    let candidate_tokens = approx_tokens(&value_text(candidate));
    for (model, note) in replies {
        let reply = note.map(reply_text).unwrap_or_default();
        if let Some(warning) =
            bias::verbosity_warning(candidate_tokens, approx_tokens(reply), threshold)
        {
            return Some(format!("juror `{model}`: {warning}"));
        }
    }
    None
}

/// Run the style-bias guard over the candidate's formatting density.
///
/// Unlike the verbosity guard, style bias is a property of the candidate alone
/// (how heavily it leans on formatting), so this needs no juror reply: it warns
/// when the candidate's density reaches the configured threshold. Returns `None`
/// when the guard is off or the candidate is below the threshold.
pub(crate) fn style_signal(bias: Option<&BiasConfig>, candidate: &Value) -> Option<String> {
    let threshold = bias?.style_guard?;
    bias::style_warning(&value_text(candidate), threshold)
}

/// Grade every anchor through every juror, then fold the deviations into a
/// [`CalibrationReport`] via [`bias::calibrate_rubric`].
///
/// This is the expensive mitigation: each anchor costs one real provider
/// call per juror, paid before the candidate is graded. Returns `None`
/// when no anchor set is configured so the free path stays free.
pub(crate) async fn run_anchor_calibration(
    config: &JuryConfig,
    providers: &[&dyn LlmProvider],
    rubric: &str,
) -> Option<CalibrationReport> {
    let anchors = config
        .bias
        .as_ref()
        .map(|b| b.anchors.as_slice())
        .filter(|a| !a.is_empty())?;
    // The math-side anchor shape carries the candidate as text; jurors see
    // the original JSON value through the judge path below.
    let math_anchors: Vec<Anchor> = anchors
        .iter()
        .map(|a| Anchor {
            response: value_text(&a.candidate),
            expected_score: a.expected,
        })
        .collect();
    let mut results = Vec::with_capacity(config.jurors.len());
    for (juror, provider) in config.jurors.iter().zip(providers.iter()) {
        results.push(
            calibrate_one_juror(config, juror, *provider, rubric, anchors, &math_anchors).await,
        );
    }
    Some(bias::fold_calibration_results(&results))
}

/// Score one juror against every anchor and fold the deviations.
///
/// The anchor calls use the same [`JudgeConfig`] the candidate call will
/// use, so a juror is calibrated on exactly the rubric it grades with. A
/// juror that produces no score for an anchor (provider error or a judge
/// malfunction) is flagged outright: it failed to grade a known-answer
/// item, which is the very miscalibration anchors exist to catch.
async fn calibrate_one_juror(
    config: &JuryConfig,
    juror: &JuryJuror,
    provider: &dyn LlmProvider,
    rubric: &str,
    anchors: &[BiasAnchor],
    math_anchors: &[Anchor],
) -> CalibrationResult {
    let judge_cfg = JudgeConfig {
        rubric: rubric.to_string(),
        threshold: juror.threshold.unwrap_or(config.threshold),
        model: Some(juror.model.clone()),
    };
    let mut observed: Vec<f64> = Vec::with_capacity(anchors.len());
    for anchor in anchors {
        let score = llm_judge::evaluate(&judge_cfg, &anchor.candidate, provider, &juror.model)
            .await
            .ok()
            .and_then(|result| extract_score(&result));
        match score {
            Some(s) => observed.push(s),
            None => {
                return CalibrationResult {
                    juror_id: juror.model.clone(),
                    miscalibrated: true,
                    deviations: Vec::new(),
                }
            }
        }
    }
    let mut scores = observed.into_iter();
    bias::calibrate_rubric(
        &juror.model,
        math_anchors,
        ANCHOR_DEVIATION_THRESHOLD,
        // `scores` is index-aligned with `math_anchors`; the fallback is
        // unreachable because every anchor pushed exactly one score above.
        |_anchor_text| scores.next().unwrap_or(0.0),
    )
}

/// Render a candidate value the way the judge prompt does, so the
/// verbosity token counts describe the text the juror actually read.
fn value_text(value: &Value) -> String {
    match value {
        Value::String(s) => s.clone(),
        other => other.to_string(),
    }
}

/// Whitespace-token approximation used by the verbosity guard.
///
/// Why not the provider's tokenizer: the guard compares a ratio, so the
/// absolute unit cancels out, and a deterministic provider-free count keeps
/// the guard from adding latency or spend to every jury run.
fn approx_tokens(text: &str) -> usize {
    text.split_whitespace().count()
}

/// Strip the bracketed provenance suffix (`[judge=... template=... ]`) the
/// judge appends to every note, so the guard measures the juror's own words
/// rather than constant bookkeeping text.
fn reply_text(note: &str) -> &str {
    if let Some(idx) = note.rfind(" [judge=") {
        return &note[..idx];
    }
    if note.starts_with("[judge=") {
        return "";
    }
    note
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    fn bias(cot: bool, guard: Option<f64>) -> BiasConfig {
        BiasConfig {
            chain_of_thought: cot,
            verbosity_guard: guard,
            style_guard: None,
            anchors: Vec::new(),
            style_swap: false,
        }
    }

    #[test]
    fn effective_rubric_appends_cot_instruction_only_when_enabled() {
        let with_cot = effective_rubric("Be kind.", Some(&bias(true, None)));
        assert!(with_cot.starts_with("Be kind."));
        assert!(with_cot.contains(CHAIN_OF_THOUGHT_INSTRUCTION));
        assert_eq!(
            effective_rubric("Be kind.", Some(&bias(false, None))),
            "Be kind."
        );
        assert_eq!(effective_rubric("Be kind.", None), "Be kind.");
    }

    #[test]
    fn verbosity_signal_ignores_the_provenance_suffix() {
        // Reply words: 3 after stripping provenance; candidate words: 4.
        // Ratio 4/3 stays under 2.0, so counting the provenance (which
        // would push the reply to 6 words and trip) must not happen.
        let note = "fine answer overall [judge=a template=v2 retried=false]";
        let replies = vec![("a", Some(note))];
        let signal = verbosity_signal(
            Some(&bias(false, Some(2.0))),
            &json!("hi there friend ok"),
            &replies,
        );
        assert!(signal.is_none(), "signal: {signal:?}");
    }

    #[test]
    fn verbosity_signal_names_the_juror_it_fired_on() {
        let long_note = format!("{} [judge=b template=v2 retried=false]", "word ".repeat(20));
        let replies = vec![("a", Some("ok fine")), ("b", Some(long_note.as_str()))];
        let signal = verbosity_signal(Some(&bias(false, Some(2.0))), &json!("two words"), &replies)
            .expect("guard trips on the 10x reply");
        assert!(signal.contains("juror `b`"), "signal: {signal}");
        assert!(signal.contains("longer response is"), "signal: {signal}");
    }

    #[test]
    fn verbosity_signal_is_silent_without_a_configured_guard() {
        let replies = vec![(
            "a",
            Some("a very long reply that would otherwise trip the guard"),
        )];
        assert!(verbosity_signal(None, &json!("hi"), &replies).is_none());
        assert!(verbosity_signal(Some(&bias(true, None)), &json!("hi"), &replies).is_none());
    }

    #[test]
    fn render_lists_active_mitigations_and_appends_signals() {
        let report = BiasReport {
            active: vec!["cot", "verbosity_guard", "style_guard"],
            verbosity_warning: Some("juror `a`: longer response is 3.00x the shorter".into()),
            style_warning: Some("candidate formatting density is 8.0 markers/100 chars".into()),
            calibration: Some(CalibrationReport {
                performed: true,
                miscalibrated_jurors: vec!["b".into()],
            }),
        };
        let rendered = report.render().expect("active mitigations render");
        assert!(
            rendered.contains("bias=[cot,verbosity_guard,style_guard]"),
            "{rendered}"
        );
        assert!(rendered.contains("verbosity: juror `a`"), "{rendered}");
        assert!(
            rendered.contains("style: candidate formatting"),
            "{rendered}"
        );
        assert!(rendered.contains("miscalibrated jurors: b"), "{rendered}");
    }

    #[test]
    fn render_is_empty_when_no_mitigation_is_configured() {
        let report = BiasReport {
            active: Vec::new(),
            verbosity_warning: None,
            style_warning: None,
            calibration: None,
        };
        assert!(report.render().is_none());
    }

    #[test]
    fn active_tags_track_each_configured_mitigation() {
        assert!(active_tags(None).is_empty());
        assert_eq!(
            active_tags(Some(&bias(true, Some(2.0)))),
            vec!["cot", "verbosity_guard"]
        );
        let with_anchors = BiasConfig {
            chain_of_thought: false,
            verbosity_guard: None,
            style_guard: None,
            anchors: vec![BiasAnchor {
                candidate: json!("known good"),
                expected: 0.9,
            }],
            style_swap: false,
        };
        assert_eq!(active_tags(Some(&with_anchors)), vec!["anchors"]);
    }

    #[test]
    fn validation_rejects_guard_ratios_at_or_below_one() {
        for bad in [1.0, 0.5, -2.0, f64::NAN] {
            let err = validate(&bias(false, Some(bad))).unwrap_err();
            assert!(err.to_string().contains("verbosity_guard"), "{err}");
        }
        assert!(validate(&bias(false, Some(1.5))).is_ok());
    }

    #[test]
    fn validation_rejects_anchor_expectations_outside_the_unit_interval() {
        for bad in [-0.1, 1.5, f64::NAN] {
            let cfg = BiasConfig {
                chain_of_thought: false,
                verbosity_guard: None,
                style_guard: None,
                anchors: vec![BiasAnchor {
                    candidate: json!("x"),
                    expected: bad,
                }],
                style_swap: false,
            };
            let err = validate(&cfg).unwrap_err();
            assert!(err.to_string().contains("anchor expected score"), "{err}");
        }
    }
}
