//! Spec-derived conformance invariants.
//!
//! The corpus `expect` DSL matches a single request/response shape. Some spec
//! requirements are properties of the whole captured session instead: the
//! handshake has to happen before anything else, a capability a server answers
//! must have been advertised, and JSON-RPC error envelopes have to carry an
//! integer code. These are not shape matches on one message, they are named
//! properties over an ordered capture.
//!
//! This module lifts those requirements into typed [`Invariant`] values, each
//! with a stable `INV-NNN` id (the family is documented in
//! `docs/conformance-invariants.md` and ADR 0015), a description, a category,
//! and a pure check function over a [`CapturedSession`]. The capture is handed
//! in already recorded so the invariant logic stays deterministic and
//! transport-free, the same contract `stream.rs` follows. Findings reuse the
//! crate's [`AssertionResult`] so reporters do not need a parallel result model.

use serde::{Deserialize, Serialize};
use serde_json::Value;

use super::invariant_checks;
use crate::matchers::AssertionResult;

/// One client-to-server exchange recorded during a session.
///
/// The request is the JSON-RPC request the client sent; the response is the
/// envelope the server returned, which carries either a `result` or an `error`.
/// A notification (no response expected) is recorded with `response` set to
/// `None`. Serde derives let the CLI load captures from a JSON file so the
/// invariant run stays deterministic and offline.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Exchange {
    /// The JSON-RPC request or notification the client sent.
    pub request: Value,
    /// The server's response envelope, or `None` for a notification.
    #[serde(default)]
    pub response: Option<Value>,
}

impl Exchange {
    /// Build an exchange from a request and its response.
    pub fn new(request: Value, response: Value) -> Self {
        Self {
            request,
            response: Some(response),
        }
    }

    /// Build an exchange for a notification the client sent (no response).
    pub fn notification(request: Value) -> Self {
        Self {
            request,
            response: None,
        }
    }

    /// The JSON-RPC `method` on the request, if any.
    pub fn method(&self) -> Option<&str> {
        self.request.get("method").and_then(Value::as_str)
    }
}

/// A captured MCP session: the ordered client exchanges plus the negotiated
/// server capabilities from `initialize`.
///
/// Recording the capture is the runner's job; the invariants only read it.
/// `server_capabilities` is the `capabilities` object from the `initialize`
/// result so capability-attestation invariants do not have to re-scan the
/// exchange list to find it. Serde derives let the CLI load a capture from a
/// JSON file.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CapturedSession {
    /// Label for the server under test, surfaced in composition reports.
    pub server_label: String,
    /// Client-to-server exchanges in the order they were sent.
    pub exchanges: Vec<Exchange>,
    /// The `capabilities` object the server returned at `initialize`, or
    /// `Value::Null` when the handshake never produced one.
    #[serde(default)]
    pub server_capabilities: Value,
}

impl CapturedSession {
    /// Build a session from a label, an exchange list, and the negotiated
    /// capability block.
    pub fn new(
        server_label: impl Into<String>,
        exchanges: Vec<Exchange>,
        server_capabilities: Value,
    ) -> Self {
        Self {
            server_label: server_label.into(),
            exchanges,
            server_capabilities,
        }
    }

    /// True when the server advertised the named top-level capability key
    /// (for example `tools`, `resources`, `prompts`) at `initialize`.
    pub fn advertised(&self, capability: &str) -> bool {
        self.server_capabilities
            .get(capability)
            .is_some_and(|v| !v.is_null())
    }

    /// All exchanges whose request `method` equals `method`.
    pub fn exchanges_for(&self, method: &str) -> Vec<&Exchange> {
        self.exchanges
            .iter()
            .filter(|e| e.method() == Some(method))
            .collect()
    }
}

/// Coarse bucket an invariant belongs to, used for grouping in reports.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum InvariantCategory {
    /// Handshake ordering and lifecycle sequencing.
    Lifecycle,
    /// Capability attestation: answering a method requires advertising it.
    Capability,
    /// Result-envelope shape for tool, resource, and prompt results.
    ResultShape,
    /// JSON-RPC error-envelope rules.
    ErrorEnvelope,
}

impl InvariantCategory {
    /// Stable lowercase label for reporter output.
    pub fn as_str(&self) -> &'static str {
        match self {
            InvariantCategory::Lifecycle => "lifecycle",
            InvariantCategory::Capability => "capability",
            InvariantCategory::ResultShape => "result-shape",
            InvariantCategory::ErrorEnvelope => "error-envelope",
        }
    }
}

/// A spec-derived invariant: a named property over a captured session.
///
/// The `check` is a pure function, so the same capture always yields the same
/// finding. `id` is a stable `INV-NNN` identifier from the family documented in
/// `docs/conformance-invariants.md`.
pub struct Invariant {
    /// Stable `INV-NNN` identifier.
    pub id: &'static str,
    /// One-line description of the property the invariant asserts.
    pub description: &'static str,
    /// Category bucket for grouping.
    pub category: InvariantCategory,
    /// Pure check over a captured session.
    pub check: fn(&CapturedSession) -> AssertionResult,
}

impl Invariant {
    /// Evaluate this invariant against a capture, tagging the result so a
    /// reporter can attribute the finding to the rule id without re-deriving
    /// it from the [`AssertionResult`].
    pub fn evaluate(&self, session: &CapturedSession) -> InvariantOutcome {
        InvariantOutcome {
            id: self.id,
            category: self.category,
            description: self.description,
            result: (self.check)(session),
        }
    }
}

/// The result of evaluating one [`Invariant`] against one capture.
///
/// Carries the rule id and category alongside the [`AssertionResult`] so the
/// composition layer and reporters can group outcomes without holding a
/// reference back to the [`Invariant`] definition.
#[derive(Debug, Clone)]
pub struct InvariantOutcome {
    /// The invariant's stable id.
    pub id: &'static str,
    /// The invariant's category.
    pub category: InvariantCategory,
    /// The invariant's description.
    pub description: &'static str,
    /// Pass/fail plus evidence.
    pub result: AssertionResult,
}

impl InvariantOutcome {
    /// True when the underlying assertion passed.
    pub fn passed(&self) -> bool {
        self.result.pass
    }
}

/// The starter set of spec-derived invariants, evaluated in id order.
///
/// Each entry is a named property, not a shape match. Callers iterate this set
/// to run every invariant against a capture.
pub const INVARIANTS: &[Invariant] = &[
    Invariant {
        id: "INV-001",
        description: "initialize is the first request the client sends",
        category: InvariantCategory::Lifecycle,
        check: invariant_checks::initialize_first,
    },
    Invariant {
        id: "INV-002",
        description: "initialized notification follows the initialize response",
        category: InvariantCategory::Lifecycle,
        check: invariant_checks::initialized_after_response,
    },
    Invariant {
        id: "INV-003",
        description: "a server answering tools/list advertised the tools capability",
        category: InvariantCategory::Capability,
        check: invariant_checks::tools_capability_attested,
    },
    Invariant {
        id: "INV-004",
        description: "an advertised capability is not erroring on every call",
        category: InvariantCategory::Capability,
        check: invariant_checks::advertised_capability_usable,
    },
    Invariant {
        id: "INV-005",
        description:
            "a tools/call result carries content or structuredContent and a boolean isError",
        category: InvariantCategory::ResultShape,
        check: invariant_checks::tool_result_shape,
    },
    Invariant {
        id: "INV-006",
        description: "JSON-RPC error envelopes carry an integer code and a string message",
        category: InvariantCategory::ErrorEnvelope,
        check: invariant_checks::error_envelope_shape,
    },
    Invariant {
        id: "INV-007",
        description: "an unknown method returns JSON-RPC error code -32601",
        category: InvariantCategory::ErrorEnvelope,
        check: invariant_checks::unknown_method_code,
    },
];

/// Run every invariant in [`INVARIANTS`] against a capture, in id order.
pub fn evaluate_all(session: &CapturedSession) -> Vec<InvariantOutcome> {
    INVARIANTS.iter().map(|inv| inv.evaluate(session)).collect()
}

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

    fn req(id: i64, method: &str) -> Value {
        json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": {} })
    }

    fn ok(id: i64, result: Value) -> Value {
        json!({ "jsonrpc": "2.0", "id": id, "result": result })
    }

    fn err(id: i64, code: i64, message: &str) -> Value {
        json!({ "jsonrpc": "2.0", "id": id, "error": { "code": code, "message": message } })
    }

    fn handshake() -> Exchange {
        Exchange::new(req(1, "initialize"), ok(1, json!({})))
    }

    fn session(exchanges: Vec<Exchange>, caps: Value) -> CapturedSession {
        CapturedSession::new("stdio://x", exchanges, caps)
    }

    fn run(id: &str, session: &CapturedSession) -> AssertionResult {
        let inv = INVARIANTS
            .iter()
            .find(|i| i.id == id)
            .expect("invariant id exists");
        (inv.check)(session)
    }

    #[test]
    fn initialize_first_passes_when_handshake_leads() {
        let s = session(
            vec![
                handshake(),
                Exchange::new(req(2, "tools/list"), ok(2, json!({}))),
            ],
            json!({ "tools": {} }),
        );
        assert!(run("INV-001", &s).pass);
    }

    #[test]
    fn initialize_first_fails_when_another_method_leads() {
        let s = session(
            vec![Exchange::new(req(1, "tools/list"), ok(1, json!({})))],
            json!({ "tools": {} }),
        );
        let r = run("INV-001", &s);
        assert!(!r.pass);
        assert!(r.diffs[0]
            .note
            .as_deref()
            .is_some_and(|n| n.contains("tools/list")));
    }

    #[test]
    fn initialize_first_fails_on_empty_capture() {
        let s = session(vec![], json!({}));
        assert!(!run("INV-001", &s).pass);
    }

    #[test]
    fn initialized_after_response_passes_in_order() {
        let s = session(
            vec![
                handshake(),
                Exchange::notification(json!({ "method": "notifications/initialized" })),
            ],
            json!({}),
        );
        assert!(run("INV-002", &s).pass);
    }

    #[test]
    fn initialized_before_response_fails() {
        let s = session(
            vec![
                Exchange::notification(json!({ "method": "notifications/initialized" })),
                handshake(),
            ],
            json!({}),
        );
        assert!(!run("INV-002", &s).pass);
    }

    #[test]
    fn initialized_without_handshake_fails() {
        let s = session(
            vec![Exchange::notification(
                json!({ "method": "notifications/initialized" }),
            )],
            json!({}),
        );
        assert!(!run("INV-002", &s).pass);
    }

    #[test]
    fn no_initialized_notification_is_trivially_ordered() {
        let s = session(vec![handshake()], json!({}));
        assert!(run("INV-002", &s).pass);
    }

    #[test]
    fn tools_list_without_advertised_capability_fails() {
        let s = session(
            vec![
                handshake(),
                Exchange::new(req(2, "tools/list"), ok(2, json!({ "tools": [] }))),
            ],
            json!({}),
        );
        let r = run("INV-003", &s);
        assert!(!r.pass);
        assert!(r.diffs[0]
            .note
            .as_deref()
            .is_some_and(|n| n.contains("tools")));
    }

    #[test]
    fn tools_list_with_advertised_capability_passes() {
        let s = session(
            vec![
                handshake(),
                Exchange::new(req(2, "tools/list"), ok(2, json!({ "tools": [] }))),
            ],
            json!({ "tools": {} }),
        );
        assert!(run("INV-003", &s).pass);
    }

    #[test]
    fn errored_list_call_does_not_require_attestation() {
        // A server that errors tools/list never "answered" it, so INV-003 holds.
        let s = session(
            vec![
                handshake(),
                Exchange::new(req(2, "tools/list"), err(2, -32601, "Method not found")),
            ],
            json!({}),
        );
        assert!(run("INV-003", &s).pass);
    }

    #[test]
    fn advertised_capability_erroring_every_call_fails() {
        let s = session(
            vec![
                handshake(),
                Exchange::new(req(2, "tools/list"), err(2, -32603, "boom")),
                Exchange::new(req(3, "tools/call"), err(3, -32603, "boom")),
            ],
            json!({ "tools": {} }),
        );
        let r = run("INV-004", &s);
        assert!(!r.pass);
        assert!(r.diffs[0]
            .note
            .as_deref()
            .is_some_and(|n| n.contains("tools")));
    }

    #[test]
    fn advertised_capability_with_one_success_passes() {
        let s = session(
            vec![
                handshake(),
                Exchange::new(req(2, "tools/list"), ok(2, json!({ "tools": [] }))),
                Exchange::new(req(3, "tools/call"), err(3, -32603, "boom")),
            ],
            json!({ "tools": {} }),
        );
        assert!(run("INV-004", &s).pass);
    }

    #[test]
    fn unadvertised_capability_is_skipped_by_inv004() {
        let s = session(
            vec![
                handshake(),
                Exchange::new(req(2, "tools/list"), err(2, -32603, "boom")),
            ],
            json!({}),
        );
        assert!(run("INV-004", &s).pass);
    }

    #[test]
    fn tool_result_with_content_array_passes() {
        let result = json!({ "content": [{ "type": "text", "text": "hi" }] });
        let s = session(
            vec![
                handshake(),
                Exchange::new(req(2, "tools/call"), ok(2, result)),
            ],
            json!({ "tools": {} }),
        );
        assert!(run("INV-005", &s).pass);
    }

    #[test]
    fn tool_result_with_structured_content_passes() {
        let result = json!({ "structuredContent": { "count": 3 } });
        let s = session(
            vec![
                handshake(),
                Exchange::new(req(2, "tools/call"), ok(2, result)),
            ],
            json!({ "tools": {} }),
        );
        assert!(run("INV-005", &s).pass);
    }

    #[test]
    fn tool_result_without_content_fails() {
        let result = json!({ "note": "missing the envelope" });
        let s = session(
            vec![
                handshake(),
                Exchange::new(req(2, "tools/call"), ok(2, result)),
            ],
            json!({ "tools": {} }),
        );
        assert!(!run("INV-005", &s).pass);
    }

    #[test]
    fn tool_result_with_non_boolean_is_error_fails() {
        let result = json!({ "content": [], "isError": "yes" });
        let s = session(
            vec![
                handshake(),
                Exchange::new(req(2, "tools/call"), ok(2, result)),
            ],
            json!({ "tools": {} }),
        );
        let r = run("INV-005", &s);
        assert!(!r.pass);
        assert!(r.diffs[0]
            .note
            .as_deref()
            .is_some_and(|n| n.contains("isError")));
    }

    #[test]
    fn errored_tool_call_skips_shape_check() {
        let s = session(
            vec![
                handshake(),
                Exchange::new(req(2, "tools/call"), err(2, -32602, "bad params")),
            ],
            json!({ "tools": {} }),
        );
        assert!(run("INV-005", &s).pass);
    }

    #[test]
    fn well_formed_error_envelope_passes() {
        let s = session(
            vec![
                handshake(),
                Exchange::new(req(2, "tools/call"), err(2, -32602, "Invalid params")),
            ],
            json!({ "tools": {} }),
        );
        assert!(run("INV-006", &s).pass);
    }

    #[test]
    fn error_envelope_with_string_code_fails() {
        let bad = json!({ "jsonrpc": "2.0", "id": 2, "error": { "code": "oops", "message": "x" } });
        let s = session(
            vec![handshake(), Exchange::new(req(2, "tools/call"), bad)],
            json!({}),
        );
        assert!(!run("INV-006", &s).pass);
    }

    #[test]
    fn error_envelope_missing_message_fails() {
        let bad = json!({ "jsonrpc": "2.0", "id": 2, "error": { "code": -32602 } });
        let s = session(
            vec![handshake(), Exchange::new(req(2, "tools/call"), bad)],
            json!({}),
        );
        assert!(!run("INV-006", &s).pass);
    }

    #[test]
    fn unknown_method_with_correct_code_passes() {
        let s = session(
            vec![
                handshake(),
                Exchange::new(req(2, "bogus/method"), err(2, -32601, "Method not found")),
            ],
            json!({}),
        );
        assert!(run("INV-007", &s).pass);
    }

    #[test]
    fn unknown_method_with_wrong_code_fails() {
        let s = session(
            vec![
                handshake(),
                Exchange::new(req(2, "bogus/method"), err(2, -32603, "Method not found")),
            ],
            json!({}),
        );
        let r = run("INV-007", &s);
        assert!(!r.pass);
        assert!(r.diffs[0]
            .note
            .as_deref()
            .is_some_and(|n| n.contains("-32601")));
    }

    #[test]
    fn unrelated_error_is_ignored_by_inv007() {
        let s = session(
            vec![
                handshake(),
                Exchange::new(req(2, "tools/call"), err(2, -32602, "Invalid params")),
            ],
            json!({}),
        );
        assert!(run("INV-007", &s).pass);
    }

    #[test]
    fn evaluate_all_runs_every_invariant_in_id_order() {
        let s = session(vec![handshake()], json!({}));
        let outcomes = evaluate_all(&s);
        assert_eq!(outcomes.len(), INVARIANTS.len());
        assert_eq!(outcomes[0].id, "INV-001");
        assert!(outcomes.iter().all(|o| o.passed()));
    }

    #[test]
    fn invariant_evaluate_tags_id_and_category() {
        let s = session(vec![handshake()], json!({}));
        let outcome = INVARIANTS[2].evaluate(&s);
        assert_eq!(outcome.id, "INV-003");
        assert_eq!(outcome.category, InvariantCategory::Capability);
    }

    #[test]
    fn category_labels_are_stable() {
        assert_eq!(InvariantCategory::Lifecycle.as_str(), "lifecycle");
        assert_eq!(InvariantCategory::Capability.as_str(), "capability");
        assert_eq!(InvariantCategory::ResultShape.as_str(), "result-shape");
        assert_eq!(InvariantCategory::ErrorEnvelope.as_str(), "error-envelope");
    }

    #[test]
    fn session_advertised_reads_capability_block() {
        let s = session(vec![handshake()], json!({ "tools": {}, "prompts": null }));
        assert!(s.advertised("tools"));
        assert!(!s.advertised("prompts"));
        assert!(!s.advertised("resources"));
    }

    #[test]
    fn exchange_notification_has_no_response() {
        let e = Exchange::notification(json!({ "method": "notifications/initialized" }));
        assert!(e.response.is_none());
        assert_eq!(e.method(), Some("notifications/initialized"));
    }

    #[test]
    fn captured_session_round_trips_through_json() {
        let original = session(
            vec![
                handshake(),
                Exchange::new(req(2, "tools/list"), ok(2, json!({ "tools": [] }))),
                Exchange::notification(json!({ "method": "notifications/initialized" })),
            ],
            json!({ "tools": {} }),
        );
        let encoded = serde_json::to_string(&original).expect("serialize");
        let decoded: CapturedSession = serde_json::from_str(&encoded).expect("deserialize");
        assert_eq!(decoded.server_label, original.server_label);
        assert_eq!(decoded.exchanges.len(), 3);
        assert!(decoded.advertised("tools"));
        // The decoded capture evaluates the same way as the original.
        assert!(evaluate_all(&decoded).iter().all(|o| o.passed()));
    }

    #[test]
    fn capture_json_defaults_missing_response_and_capabilities() {
        // A notification-only exchange omits `response`; a session may omit
        // `server_capabilities`. Both default cleanly.
        let wire = json!({
            "server_label": "x",
            "exchanges": [
                { "request": { "jsonrpc": "2.0", "id": 1, "method": "initialize" },
                  "response": { "jsonrpc": "2.0", "id": 1, "result": {} } },
                { "request": { "method": "notifications/initialized" } }
            ]
        });
        let decoded: CapturedSession = serde_json::from_value(wire).expect("deserialize");
        assert_eq!(decoded.exchanges.len(), 2);
        assert!(decoded.exchanges[1].response.is_none());
        assert!(decoded.server_capabilities.is_null());
    }

    #[test]
    fn invariant_ids_are_unique() {
        let mut ids: Vec<&str> = INVARIANTS.iter().map(|i| i.id).collect();
        ids.sort_unstable();
        ids.dedup();
        assert_eq!(ids.len(), INVARIANTS.len(), "duplicate INV id detected");
    }
}
