//! Model-agnostic invariant checks.
//!
//! Invariants are the guardrails a workflow needs regardless of which
//! model is in the loop. The diff engine integrates these results into
//! [`crate::model_compat::BaselineDiff`] so an invariant failure
//! classifies as `FAIL` even when the surface diff alone would have
//! classified as `DRIFT`.

use serde::{Deserialize, Serialize};

use crate::model_compat::model::{CapturedRun, ContentBlock, FinishReason};

/// One named invariant declared on a test or suite.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Invariant {
    /// User-facing name. Used in reports and as the violation id.
    pub name: String,
    /// What to evaluate.
    pub kind: InvariantKind,
    /// Payload for the chosen kind.
    pub condition: InvariantCondition,
}

/// Discriminant for [`Invariant`]. Each value matches a payload variant
/// of [`InvariantCondition`].
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum InvariantKind {
    /// A specific tool must appear in `tool_calls`.
    ToolCalled,
    /// A tool argument must be present (no value constraint).
    ArgPresent,
    /// A response content block must mention a substring.
    ResponseFieldPresent,
    /// A response must semantically match a reference (case-insensitive
    /// substring; the live engine would defer to the jury or embeddings,
    /// but the OSS guardrail is a deterministic check).
    ResponseSemanticMatch,
    /// Observed latency must be below the configured ceiling.
    LatencyUnderMs,
}

/// Payload for an invariant. The shape mirrors the `kind` field.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum InvariantCondition {
    /// Single tool name (for [`InvariantKind::ToolCalled`]).
    Tool {
        /// Name of the tool that must be called.
        tool: String,
    },
    /// Tool plus argument name (for [`InvariantKind::ArgPresent`]).
    Arg {
        /// Owning tool.
        tool: String,
        /// Required argument key.
        arg: String,
    },
    /// Free-form text snippet (for [`InvariantKind::ResponseFieldPresent`]
    /// and [`InvariantKind::ResponseSemanticMatch`]).
    Text {
        /// Substring or reference text.
        text: String,
    },
    /// Latency ceiling in milliseconds (for
    /// [`InvariantKind::LatencyUnderMs`]).
    Latency {
        /// Inclusive ceiling.
        ms: u64,
    },
}

/// Outcome of evaluating one invariant against one captured run.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct InvariantResult {
    /// Name of the invariant that ran.
    pub name: String,
    /// True when the invariant was satisfied.
    pub passed: bool,
    /// Short, human-readable reason. Populated even on pass so reports
    /// can show "tool `lookup_account` called as required".
    pub detail: String,
}

/// Evaluate every invariant against the captured run and return one
/// result per invariant. Order matches the input.
///
/// Invariant failures classify as `FAIL` in the diff layer, never as
/// `DRIFT`. The diff engine reads this list and surfaces each failure
/// alongside the surface-level diff entries.
pub fn check_invariants(invariants: &[Invariant], run: &CapturedRun) -> Vec<InvariantResult> {
    invariants
        .iter()
        .map(|inv| evaluate_one(inv, run))
        .collect()
}

fn evaluate_one(inv: &Invariant, run: &CapturedRun) -> InvariantResult {
    match (inv.kind, &inv.condition) {
        (InvariantKind::ToolCalled, InvariantCondition::Tool { tool }) => {
            check_tool_called(&inv.name, tool, run)
        }
        (InvariantKind::ArgPresent, InvariantCondition::Arg { tool, arg }) => {
            check_arg_present(&inv.name, tool, arg, run)
        }
        (InvariantKind::ResponseFieldPresent, InvariantCondition::Text { text }) => {
            check_response_field_present(&inv.name, text, run)
        }
        (InvariantKind::ResponseSemanticMatch, InvariantCondition::Text { text }) => {
            check_response_semantic_match(&inv.name, text, run)
        }
        (InvariantKind::LatencyUnderMs, InvariantCondition::Latency { ms }) => {
            check_latency_under(&inv.name, *ms, run)
        }
        _ => InvariantResult {
            name: inv.name.clone(),
            passed: false,
            detail: "invariant kind/condition mismatch".into(),
        },
    }
}

fn check_tool_called(name: &str, tool: &str, run: &CapturedRun) -> InvariantResult {
    let called = run.exchange.tool_calls.iter().any(|c| c.name == tool);
    InvariantResult {
        name: name.into(),
        passed: called,
        detail: if called {
            format!("tool `{tool}` called as required")
        } else {
            format!("required tool `{tool}` not called")
        },
    }
}

fn check_arg_present(name: &str, tool: &str, arg: &str, run: &CapturedRun) -> InvariantResult {
    let found = run
        .exchange
        .tool_calls
        .iter()
        .filter(|c| c.name == tool)
        .any(|c| c.arguments.get(arg).is_some());
    InvariantResult {
        name: name.into(),
        passed: found,
        detail: if found {
            format!("arg `{arg}` present on `{tool}`")
        } else {
            format!("required arg `{arg}` missing on `{tool}`")
        },
    }
}

fn check_response_field_present(name: &str, needle: &str, run: &CapturedRun) -> InvariantResult {
    let found = run.exchange.content.iter().any(|b| match b {
        ContentBlock::Text { text } => text.contains(needle),
        ContentBlock::Other => false,
    });
    InvariantResult {
        name: name.into(),
        passed: found,
        detail: if found {
            format!("response contains `{needle}`")
        } else {
            format!("response missing required substring `{needle}`")
        },
    }
}

fn check_response_semantic_match(
    name: &str,
    reference: &str,
    run: &CapturedRun,
) -> InvariantResult {
    let needle = reference.to_lowercase();
    let found = run.exchange.content.iter().any(|b| match b {
        ContentBlock::Text { text } => text.to_lowercase().contains(&needle),
        ContentBlock::Other => false,
    });
    InvariantResult {
        name: name.into(),
        passed: found,
        detail: if found {
            format!("response semantically matches `{reference}`")
        } else {
            format!("response does not semantically match `{reference}`")
        },
    }
}

fn check_latency_under(name: &str, ceiling: u64, run: &CapturedRun) -> InvariantResult {
    let ok = run.latency_ms <= ceiling;
    InvariantResult {
        name: name.into(),
        passed: ok,
        detail: if ok {
            format!("latency {}ms under ceiling {}ms", run.latency_ms, ceiling)
        } else {
            format!(
                "latency {}ms exceeded ceiling {}ms",
                run.latency_ms, ceiling
            )
        },
    }
}

/// Note for the finish_reason discriminant. Some tests prefer to pin
/// finish_reason via an invariant so the rollout fails fast when a
/// model flips from `stop` to `max_tokens`. The check is included as a
/// helper free function rather than a kind because we keep the
/// invariant taxonomy narrow per ADR 0026.
pub fn finish_reason_equals(run: &CapturedRun, expected: FinishReason) -> bool {
    run.exchange.finish_reason == expected
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::model_compat::model::{CapturedExchange, ToolCall};

    fn run_with_tools(tools: Vec<ToolCall>) -> CapturedRun {
        CapturedRun::new(
            CapturedExchange {
                id: "t::a::1".into(),
                content: vec![ContentBlock::Text {
                    text: "Account lookup complete.".into(),
                }],
                tool_calls: tools,
                finish_reason: FinishReason::ToolUse,
                mcp_messages: vec![],
            },
            120,
        )
    }

    #[test]
    fn tool_called_passes_when_present() {
        let inv = Invariant {
            name: "must-call-lookup".into(),
            kind: InvariantKind::ToolCalled,
            condition: InvariantCondition::Tool {
                tool: "lookup_account".into(),
            },
        };
        let run = run_with_tools(vec![ToolCall {
            name: "lookup_account".into(),
            arguments: serde_json::json!({"id": "1"}),
            raw_arguments: None,
        }]);
        let r = check_invariants(&[inv], &run);
        assert!(r[0].passed);
    }

    #[test]
    fn tool_called_fails_when_absent() {
        let inv = Invariant {
            name: "must-call-lookup".into(),
            kind: InvariantKind::ToolCalled,
            condition: InvariantCondition::Tool {
                tool: "lookup_account".into(),
            },
        };
        let run = run_with_tools(vec![]);
        let r = check_invariants(&[inv], &run);
        assert!(!r[0].passed);
    }

    #[test]
    fn arg_present_passes_when_key_exists() {
        let inv = Invariant {
            name: "lookup-needs-id".into(),
            kind: InvariantKind::ArgPresent,
            condition: InvariantCondition::Arg {
                tool: "lookup_account".into(),
                arg: "id".into(),
            },
        };
        let run = run_with_tools(vec![ToolCall {
            name: "lookup_account".into(),
            arguments: serde_json::json!({"id": "1"}),
            raw_arguments: None,
        }]);
        let r = check_invariants(&[inv], &run);
        assert!(r[0].passed);
    }

    #[test]
    fn arg_present_fails_when_key_missing() {
        let inv = Invariant {
            name: "lookup-needs-id".into(),
            kind: InvariantKind::ArgPresent,
            condition: InvariantCondition::Arg {
                tool: "lookup_account".into(),
                arg: "id".into(),
            },
        };
        let run = run_with_tools(vec![ToolCall {
            name: "lookup_account".into(),
            arguments: serde_json::json!({}),
            raw_arguments: None,
        }]);
        let r = check_invariants(&[inv], &run);
        assert!(!r[0].passed);
    }

    #[test]
    fn response_field_present_substring_match() {
        let inv = Invariant {
            name: "mentions-account".into(),
            kind: InvariantKind::ResponseFieldPresent,
            condition: InvariantCondition::Text {
                text: "Account".into(),
            },
        };
        let run = run_with_tools(vec![]);
        let r = check_invariants(&[inv], &run);
        assert!(r[0].passed);
    }

    #[test]
    fn response_semantic_match_is_case_insensitive() {
        let inv = Invariant {
            name: "mentions-account".into(),
            kind: InvariantKind::ResponseSemanticMatch,
            condition: InvariantCondition::Text {
                text: "ACCOUNT LOOKUP".into(),
            },
        };
        let run = run_with_tools(vec![]);
        let r = check_invariants(&[inv], &run);
        assert!(r[0].passed);
    }

    #[test]
    fn latency_under_ms_passes_under_ceiling() {
        let inv = Invariant {
            name: "p95-budget".into(),
            kind: InvariantKind::LatencyUnderMs,
            condition: InvariantCondition::Latency { ms: 200 },
        };
        let run = run_with_tools(vec![]);
        let r = check_invariants(&[inv], &run);
        assert!(r[0].passed);
    }

    #[test]
    fn latency_under_ms_fails_over_ceiling() {
        let inv = Invariant {
            name: "p95-budget".into(),
            kind: InvariantKind::LatencyUnderMs,
            condition: InvariantCondition::Latency { ms: 10 },
        };
        let run = run_with_tools(vec![]);
        let r = check_invariants(&[inv], &run);
        assert!(!r[0].passed);
    }

    #[test]
    fn mismatched_kind_and_condition_fails() {
        let inv = Invariant {
            name: "broken".into(),
            kind: InvariantKind::ToolCalled,
            condition: InvariantCondition::Latency { ms: 100 },
        };
        let run = run_with_tools(vec![]);
        let r = check_invariants(&[inv], &run);
        assert!(!r[0].passed);
    }
}
