//! Data model for captured baselines (ADR 0026).
//!
//! These types are the on-disk shape of `baselines/<model>.json` plus
//! the in-memory `CapturedRun` the diff engine consumes. They are
//! transport-free: a real runner populates them from live MCP traffic
//! and tests populate them by hand.

use std::collections::BTreeMap;

use serde::{Deserialize, Serialize};

/// Schema URL pinned in every baseline file. Mirrors the URL pattern
/// from ADR 0026 so consumers can validate against the published JSON
/// Schema.
pub const BASELINE_SCHEMA_URL: &str = "https://mcptest.sh/schema/model-compat/v1.json";

/// Schema major version. Bumped on breaking changes per ADR 0001.
pub const BASELINE_SCHEMA_VERSION: u32 = 1;

/// On-disk baseline file. One file per (model, suite) pair.
///
/// The five identity fields (`model_id`, `model_version`,
/// `mcptest_version`, `server_version`, `server_config_hash`) pin the
/// snapshot to a specific moment so a re-capture with the same inputs
/// produces a byte-identical file (ADR 0026 idempotency contract).
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct BaselineFile {
    /// JSON Schema URL the file validates against.
    pub schema: String,
    /// Schema major version. Diff refuses cross-version comparisons.
    pub version: u32,
    /// Provider-qualified model identifier (for example
    /// `anthropic:claude-sonnet-4.5`).
    pub model_id: String,
    /// Model version string as reported by the provider. Distinct from
    /// `model_id` because two snapshots of `claude-sonnet-4.5` taken
    /// weeks apart can carry different point releases.
    pub model_version: String,
    /// Version of the `mcptest` binary that produced the capture.
    pub mcptest_version: String,
    /// Version of the MCP server under test, as reported by the
    /// `initialize` handshake.
    pub server_version: String,
    /// SHA-256 over the resolved server configuration so a config
    /// change invalidates the baseline.
    pub server_config_hash: String,
    /// Invariant and variance rule lists active at capture time.
    pub config_fingerprint: ConfigFingerprint,
    /// Per-test captured exchanges, sorted by `id` for determinism.
    pub assertions: Vec<CapturedExchange>,
    /// Invariant declarations the diff engine should enforce. Empty
    /// when the suite did not author any (ADR 0026 default still
    /// applies via derived invariants).
    #[serde(default)]
    pub invariants: Vec<crate::model_compat::invariants::Invariant>,
}

/// The rule lists that were live when the baseline was captured.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ConfigFingerprint {
    /// Canonical invariant names (`tool_called.*`, `tool_args.*`, etc.).
    pub invariants: Vec<String>,
    /// Canonical variance names (`text_content`, `field_order`, etc.).
    pub variances: Vec<String>,
}

impl Default for ConfigFingerprint {
    fn default() -> Self {
        Self {
            invariants: default_invariants(),
            variances: default_variances(),
        }
    }
}

/// Default invariants per ADR 0026. Users can restrict the list per
/// suite; the engine never goes below these unless the user opts out.
pub fn default_invariants() -> Vec<String> {
    vec![
        "tool_called".into(),
        "tool_order".into(),
        "tool_args".into(),
        "response_shape".into(),
        "finish_reason".into(),
    ]
}

/// Default variances per ADR 0026.
pub fn default_variances() -> Vec<String> {
    vec![
        "text_content".into(),
        "field_order".into(),
        "whitespace".into(),
        "case".into(),
        "additive_fields".into(),
    ]
}

/// One captured assertion, keyed by stable id (`file::test::step`).
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CapturedExchange {
    /// Stable identifier of the form `<file>::<test-name>::<step>`.
    pub id: String,
    /// Response content blocks (text, image, etc.).
    pub content: Vec<ContentBlock>,
    /// Ordered list of tool calls the model emitted.
    pub tool_calls: Vec<ToolCall>,
    /// Finish reason reported by the provider.
    pub finish_reason: FinishReason,
    /// Raw MCP messages exchanged for this assertion. Useful for
    /// auditing what the runner observed when the verdict is ambiguous.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub mcp_messages: Vec<McpMessage>,
}

/// A single content block from the response. We model only `text` for
/// v1 (other block types pass through as JSON) so the diff engine has
/// one shape to walk.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ContentBlock {
    /// Plain-text content. Subject to the `text_content` variance rule.
    Text {
        /// The text body.
        text: String,
    },
    /// Any other block type passes through verbatim. Captured for
    /// debuggability; the diff engine compares structurally.
    #[serde(other)]
    Other,
}

/// A tool call the model emitted. `arguments` is a JSON object; the
/// diff engine normalizes key order when comparing.
///
/// `raw_arguments` exists so the diff engine can distinguish two bags
/// that parse to byte-equal [`serde_json::Value`]s but differ at the
/// source level (key reorder). The default `Map = BTreeMap` backing
/// `Value` discards source key order, so without raw bytes the engine
/// misclassifies key reorder as PASS instead of `DRIFT`
/// (`structurally-equivalent`). Callers that do not care leave the
/// field `None`; the diff falls back to `Value` equality.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolCall {
    /// Tool name as registered by the MCP server.
    pub name: String,
    /// Argument bag, captured as JSON so the diff engine can walk into
    /// nested structures without losing types.
    pub arguments: serde_json::Value,
    /// Original JSON source bytes of `arguments`. Used by the diff
    /// engine for key-reorder drift detection. Skipped on serialize
    /// when `None` so existing baselines round-trip byte-identical.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub raw_arguments: Option<String>,
}

/// Provider-agnostic finish reason. ADR 0026 makes this an invariant by
/// default; a server cannot silently flip `stop` to `max_tokens`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum FinishReason {
    /// Model decided it was done.
    Stop,
    /// Model emitted a tool call.
    ToolUse,
    /// Response truncated by the model's token budget.
    MaxTokens,
    /// Provider returned a finish reason the engine does not model.
    /// Captured verbatim so the diff layer can compare strings.
    Other,
}

/// A raw MCP message captured during the run. Stored as JSON so the
/// engine does not have to parse the protocol envelope to round-trip.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct McpMessage {
    /// Direction of travel: `client` or `server`.
    pub direction: String,
    /// JSON-RPC method or notification name.
    pub method: String,
    /// Full message body as captured.
    pub body: serde_json::Value,
}

/// In-memory view of a single captured assertion plus the surrounding
/// context the invariant engine needs (latency, raw exchange).
#[derive(Debug, Clone)]
pub struct CapturedRun {
    /// The captured exchange.
    pub exchange: CapturedExchange,
    /// Observed wall-clock latency in milliseconds.
    pub latency_ms: u64,
}

impl CapturedRun {
    /// Convenience constructor for tests and the runner.
    pub fn new(exchange: CapturedExchange, latency_ms: u64) -> Self {
        Self {
            exchange,
            latency_ms,
        }
    }
}

/// Helper that converts a list of [`CapturedExchange`] into a map keyed
/// by stable id. Used by the diff engine when pairing baseline and
/// candidate.
pub(crate) fn index_by_id(items: &[CapturedExchange]) -> BTreeMap<&str, &CapturedExchange> {
    let mut out = BTreeMap::new();
    for ex in items {
        out.insert(ex.id.as_str(), ex);
    }
    out
}

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

    #[test]
    fn defaults_match_adr_0026() {
        let inv = default_invariants();
        let var = default_variances();
        assert!(inv.contains(&"tool_called".to_string()));
        assert!(inv.contains(&"tool_order".to_string()));
        assert!(inv.contains(&"tool_args".to_string()));
        assert!(inv.contains(&"response_shape".to_string()));
        assert!(inv.contains(&"finish_reason".to_string()));
        assert_eq!(var.len(), 5);
    }

    #[test]
    fn content_block_text_round_trips() {
        let b = ContentBlock::Text {
            text: "hello".into(),
        };
        let s = serde_json::to_string(&b).expect("serialize");
        let parsed: ContentBlock = serde_json::from_str(&s).expect("parse");
        assert_eq!(b, parsed);
    }

    #[test]
    fn baseline_file_round_trips() {
        let bf = BaselineFile {
            schema: BASELINE_SCHEMA_URL.into(),
            version: BASELINE_SCHEMA_VERSION,
            model_id: "anthropic:claude-sonnet-4.5".into(),
            model_version: "20251022".into(),
            mcptest_version: "0.1.0".into(),
            server_version: "1.0.0".into(),
            server_config_hash: "sha256:0".into(),
            config_fingerprint: ConfigFingerprint::default(),
            assertions: vec![],
            invariants: vec![],
        };
        let s = serde_json::to_string(&bf).expect("serialize");
        let parsed: BaselineFile = serde_json::from_str(&s).expect("parse");
        assert_eq!(bf, parsed);
    }
}
