//! Multi-server composition-safety checks.
//!
//! A single MCP client increasingly fronts several servers at once. Each server
//! can be perfectly conformant on its own and still create a hazard when it
//! coexists with another behind one client: two servers exposing the same tool
//! name collide, and two servers sharing a transport can collide on the
//! JSON-RPC id space or on notification routing. Those are properties of the
//! pair, not of either capture alone.
//!
//! This module runs each server's [`Invariant`](super::invariants::Invariant)
//! set individually, then
//! re-asserts the cross-server properties that only fail when servers coexist.
//! It is pure over the captured inputs, so it stays deterministic and
//! CLI-friendly. The cross-server aggregation and governance layer (fleet
//! dashboards, org policy) is enterprise scope and is not here:
//! this is per-run, single-developer composition safety only.

use std::collections::BTreeMap;

use serde_json::Value;

use super::invariants::{evaluate_all, CapturedSession, InvariantOutcome};

/// The per-server invariant run feeding a composition report.
///
/// Holds the server label and the outcome of every invariant evaluated against
/// that server's capture, so a reporter can show single-server findings next to
/// the composition-level ones.
#[derive(Debug, Clone)]
pub struct ServerInvariantRun {
    /// Label of the server under test.
    pub server_label: String,
    /// Outcome of every invariant for this server, in id order.
    pub outcomes: Vec<InvariantOutcome>,
}

impl ServerInvariantRun {
    /// True when every invariant for this server passed.
    pub fn all_passed(&self) -> bool {
        self.outcomes.iter().all(InvariantOutcome::passed)
    }
}

/// Which composition hazard a [`CompositionFinding`] reports.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CompositionKind {
    /// Two servers expose the same unprefixed tool name.
    ToolNamespaceOverlap,
    /// Two servers sharing a transport reuse the same JSON-RPC id.
    SharedTransportIdCollision,
}

impl CompositionKind {
    /// Stable lowercase label for reporter output.
    pub fn as_str(&self) -> &'static str {
        match self {
            CompositionKind::ToolNamespaceOverlap => "tool-namespace-overlap",
            CompositionKind::SharedTransportIdCollision => "shared-transport-id-collision",
        }
    }
}

/// One composition-level finding across two or more servers.
///
/// Composition findings are always failures: a passing composition produces no
/// findings. The `detail` is reporter-ready and names the colliding servers and
/// the colliding key.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CompositionFinding {
    /// The hazard class.
    pub kind: CompositionKind,
    /// Reporter-ready description naming the servers and colliding key.
    pub detail: String,
    /// Server labels involved in the collision, sorted for determinism.
    pub servers: Vec<String>,
}

/// The full composition-safety result: per-server invariant runs plus the
/// composition-level findings.
///
/// `composition_findings` is empty when no cross-server hazard was found. A
/// caller decides the exit code with [`Self::all_safe`].
#[derive(Debug, Clone)]
pub struct CompositionReport {
    /// Per-server invariant outcomes.
    pub server_runs: Vec<ServerInvariantRun>,
    /// Cross-server hazards. Empty when the composition is safe.
    pub composition_findings: Vec<CompositionFinding>,
}

impl CompositionReport {
    /// True when every server's invariants passed and no composition hazard was
    /// found. This is the CI-green condition.
    pub fn all_safe(&self) -> bool {
        self.server_runs.iter().all(ServerInvariantRun::all_passed)
            && self.composition_findings.is_empty()
    }

    /// Count of failing single-server invariants across every server.
    pub fn failed_invariant_count(&self) -> usize {
        self.server_runs
            .iter()
            .flat_map(|r| r.outcomes.iter())
            .filter(|o| !o.passed())
            .count()
    }
}

/// Run each server's invariants, then re-assert the cross-server composition
/// properties. Pure over the captured sessions.
///
/// ```
/// use serde_json::json;
/// use mcptest_core::compliance::{check_composition, CapturedSession, Exchange};
///
/// fn tools_list(names: &[&str]) -> Exchange {
///     let tools: Vec<_> = names.iter().map(|n| json!({ "name": n })).collect();
///     Exchange::new(
///         json!({ "jsonrpc": "2.0", "id": 2, "method": "tools/list" }),
///         json!({ "jsonrpc": "2.0", "id": 2, "result": { "tools": tools } }),
///     )
/// }
/// fn handshake() -> Exchange {
///     Exchange::new(
///         json!({ "jsonrpc": "2.0", "id": 1, "method": "initialize" }),
///         json!({ "jsonrpc": "2.0", "id": 1, "result": {} }),
///     )
/// }
/// let a = CapturedSession::new("a", vec![handshake(), tools_list(&["search"])], json!({ "tools": {} }));
/// let b = CapturedSession::new("b", vec![handshake(), tools_list(&["search"])], json!({ "tools": {} }));
/// let report = check_composition(&[a, b]);
/// assert!(!report.all_safe());
/// ```
pub fn check_composition(sessions: &[CapturedSession]) -> CompositionReport {
    let server_runs = sessions
        .iter()
        .map(|s| ServerInvariantRun {
            server_label: s.server_label.clone(),
            outcomes: evaluate_all(s),
        })
        .collect();
    let mut composition_findings = check_tool_namespace_overlap(sessions);
    composition_findings.extend(check_shared_transport_id_collision(sessions));
    CompositionReport {
        server_runs,
        composition_findings,
    }
}

/// Collect the tool names a server advertised across every `tools/list` result.
fn advertised_tool_names(session: &CapturedSession) -> Vec<String> {
    let mut names = Vec::new();
    for exchange in session.exchanges_for("tools/list") {
        let Some(response) = exchange.response.as_ref() else {
            continue;
        };
        let Some(tools) = response.pointer("/result/tools").and_then(Value::as_array) else {
            continue;
        };
        for tool in tools {
            if let Some(name) = tool.get("name").and_then(Value::as_str) {
                names.push(name.to_string());
            }
        }
    }
    names
}

/// Tool-namespace-overlap: two servers exposing the same unprefixed tool name
/// under one client are a hazard. The driver namespaces tools as
/// `<server>__tool`, so an unprefixed collision means a client that forgot to
/// namespace would route ambiguously. Reports the offending name and the
/// servers that share it.
fn check_tool_namespace_overlap(sessions: &[CapturedSession]) -> Vec<CompositionFinding> {
    let mut owners: BTreeMap<String, Vec<String>> = BTreeMap::new();
    for session in sessions {
        for name in advertised_tool_names(session) {
            let entry = owners.entry(name).or_default();
            if !entry.contains(&session.server_label) {
                entry.push(session.server_label.clone());
            }
        }
    }
    owners
        .into_iter()
        .filter(|(_, servers)| servers.len() > 1)
        .map(|(name, mut servers)| {
            servers.sort();
            CompositionFinding {
                kind: CompositionKind::ToolNamespaceOverlap,
                detail: format!(
                    "tool name `{name}` is exposed by {} servers ({}); namespace it as <server>__{name}",
                    servers.len(),
                    servers.join(", ")
                ),
                servers,
            }
        })
        .collect()
}

/// Collect the JSON-RPC request ids a server used across its exchanges.
fn used_request_ids(session: &CapturedSession) -> Vec<String> {
    session
        .exchanges
        .iter()
        .filter_map(|e| e.request.get("id"))
        .map(canonical_id)
        .collect()
}

/// Render a JSON-RPC id (number or string) into a comparable canonical form so
/// `1` and `"1"` are treated as the same slot, which is what a shared transport
/// would experience.
fn canonical_id(id: &Value) -> String {
    match id {
        Value::Number(n) => n.to_string(),
        Value::String(s) => s.clone(),
        other => other.to_string(),
    }
}

/// Shared-transport-id-collision: when two servers share one transport, the
/// client multiplexes both id spaces over a single channel, so the same id used
/// by both servers cannot be correlated back to one request. A client that
/// gives each server its own transport avoids this; this check flags the
/// overlap so the operator knows a shared-transport setup would interfere.
fn check_shared_transport_id_collision(sessions: &[CapturedSession]) -> Vec<CompositionFinding> {
    let mut findings = Vec::new();
    for (i, a) in sessions.iter().enumerate() {
        for b in sessions.iter().skip(i + 1) {
            let a_ids = used_request_ids(a);
            let shared: Vec<String> = used_request_ids(b)
                .into_iter()
                .filter(|id| a_ids.contains(id))
                .collect();
            if !shared.is_empty() {
                let mut servers = vec![a.server_label.clone(), b.server_label.clone()];
                servers.sort();
                findings.push(CompositionFinding {
                    kind: CompositionKind::SharedTransportIdCollision,
                    detail: format!(
                        "servers {} and {} both use request id(s) [{}]; on a shared transport these cannot be correlated",
                        servers[0],
                        servers[1],
                        shared.join(", ")
                    ),
                    servers,
                });
            }
        }
    }
    findings
}

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

    fn handshake(id: i64) -> Exchange {
        Exchange::new(
            json!({ "jsonrpc": "2.0", "id": id, "method": "initialize" }),
            json!({ "jsonrpc": "2.0", "id": id, "result": {} }),
        )
    }

    fn tools_list(id: i64, names: &[&str]) -> Exchange {
        let tools: Vec<Value> = names.iter().map(|n| json!({ "name": n })).collect();
        Exchange::new(
            json!({ "jsonrpc": "2.0", "id": id, "method": "tools/list" }),
            json!({ "jsonrpc": "2.0", "id": id, "result": { "tools": tools } }),
        )
    }

    fn server(label: &str, exchanges: Vec<Exchange>) -> CapturedSession {
        CapturedSession::new(label, exchanges, json!({ "tools": {} }))
    }

    #[test]
    fn two_servers_without_overlap_are_safe() {
        let a = server(
            "alpha",
            vec![handshake(1), tools_list(2, &["alpha_search"])],
        );
        let b = server("beta", vec![handshake(3), tools_list(4, &["beta_lookup"])]);
        let report = check_composition(&[a, b]);
        assert!(
            report.all_safe(),
            "findings: {:?}",
            report.composition_findings
        );
    }

    #[test]
    fn overlapping_tool_name_is_flagged() {
        let a = server("alpha", vec![handshake(1), tools_list(2, &["search"])]);
        let b = server("beta", vec![handshake(3), tools_list(4, &["search"])]);
        let report = check_composition(&[a, b]);
        assert!(!report.all_safe());
        let overlap = report
            .composition_findings
            .iter()
            .find(|f| f.kind == CompositionKind::ToolNamespaceOverlap)
            .expect("overlap reported");
        assert!(overlap.detail.contains("search"));
        assert_eq!(
            overlap.servers,
            vec!["alpha".to_string(), "beta".to_string()]
        );
    }

    #[test]
    fn overlap_lists_servers_sorted_and_deduped() {
        // alpha lists `search` twice; it should not appear as its own collision.
        let a = server(
            "alpha",
            vec![
                handshake(1),
                tools_list(2, &["search"]),
                tools_list(3, &["search"]),
            ],
        );
        let report = check_composition(&[a]);
        assert!(
            report.composition_findings.is_empty(),
            "a single server listing a name twice is not a cross-server collision"
        );
    }

    #[test]
    fn shared_request_id_across_servers_is_flagged() {
        // Both servers reuse id 2 for their tools/list, which collides on a
        // shared transport.
        let a = server("alpha", vec![handshake(1), tools_list(2, &["a_tool"])]);
        let b = server("beta", vec![handshake(1), tools_list(2, &["b_tool"])]);
        let report = check_composition(&[a, b]);
        let collision = report
            .composition_findings
            .iter()
            .find(|f| f.kind == CompositionKind::SharedTransportIdCollision)
            .expect("id collision reported");
        assert!(collision.detail.contains("alpha"));
        assert!(collision.detail.contains("beta"));
    }

    #[test]
    fn disjoint_request_ids_do_not_collide() {
        let a = server("alpha", vec![handshake(1), tools_list(2, &["a_tool"])]);
        let b = server("beta", vec![handshake(3), tools_list(4, &["b_tool"])]);
        let report = check_composition(&[a, b]);
        assert!(report
            .composition_findings
            .iter()
            .all(|f| f.kind != CompositionKind::SharedTransportIdCollision));
    }

    #[test]
    fn numeric_and_string_ids_collide_on_canonical_form() {
        let a = server(
            "alpha",
            vec![Exchange::new(
                json!({ "jsonrpc": "2.0", "id": 7, "method": "initialize" }),
                json!({ "jsonrpc": "2.0", "id": 7, "result": {} }),
            )],
        );
        let b = server(
            "beta",
            vec![Exchange::new(
                json!({ "jsonrpc": "2.0", "id": "7", "method": "initialize" }),
                json!({ "jsonrpc": "2.0", "id": "7", "result": {} }),
            )],
        );
        let report = check_composition(&[a, b]);
        assert!(report
            .composition_findings
            .iter()
            .any(|f| f.kind == CompositionKind::SharedTransportIdCollision));
    }

    #[test]
    fn per_server_invariant_failure_surfaces_in_report() {
        // beta answers tools/list but did not advertise tools, so INV-003 fails
        // for beta even though the composition itself is clean.
        let a = server("alpha", vec![handshake(1), tools_list(2, &["a_tool"])]);
        let b = CapturedSession::new(
            "beta",
            vec![handshake(3), tools_list(4, &["b_tool"])],
            json!({}),
        );
        let report = check_composition(&[a, b]);
        assert!(!report.all_safe());
        assert!(report.failed_invariant_count() >= 1);
        let beta_run = report
            .server_runs
            .iter()
            .find(|r| r.server_label == "beta")
            .expect("beta run present");
        assert!(!beta_run.all_passed());
    }

    #[test]
    fn empty_input_is_trivially_safe() {
        let report = check_composition(&[]);
        assert!(report.all_safe());
        assert_eq!(report.failed_invariant_count(), 0);
    }

    #[test]
    fn three_way_overlap_names_all_servers() {
        let a = server("a", vec![handshake(1), tools_list(2, &["dup"])]);
        let b = server("b", vec![handshake(3), tools_list(4, &["dup"])]);
        let c = server("c", vec![handshake(5), tools_list(6, &["dup"])]);
        let report = check_composition(&[a, b, c]);
        let overlap = report
            .composition_findings
            .iter()
            .find(|f| f.kind == CompositionKind::ToolNamespaceOverlap)
            .expect("overlap reported");
        assert_eq!(overlap.servers.len(), 3);
    }

    #[test]
    fn composition_kind_labels_are_stable() {
        assert_eq!(
            CompositionKind::ToolNamespaceOverlap.as_str(),
            "tool-namespace-overlap"
        );
        assert_eq!(
            CompositionKind::SharedTransportIdCollision.as_str(),
            "shared-transport-id-collision"
        );
    }
}
