//! Live capture for `mcptest conformance run --server` (WOR-1382).
//!
//! Connects to the server under test, runs the MCP handshake, and records the
//! observable facts the universal conformance probe needs
//! ([`ConformanceEvidence`]). The capture is deliberately thin: it gathers
//! evidence and hands it to the pure analyzer in
//! [`mcptest_core::conformance::probe`], so the verdict logic stays
//! offline-testable and the live step is the only part that touches a socket.

use std::collections::BTreeMap;

use anyhow::{Context, Result};
use mcptest_core::conformance::{
    CancellationObservation, ConformanceEvidence, ResourceReadObservation,
};
use mcptest_core::connector::{connect_server, ConnectError};
use mcptest_core::network::ProxyConfig;
use mcptest_core::protocol::{Client, ProtocolError};
use mcptest_core::suite::ServerSpec;
use serde_json::{json, Value};

/// A URI no real server should resolve, used to probe the SEP-2164
/// resource-not-found behavior.
const BOGUS_RESOURCE_URI: &str = "mcptest://conformance-probe/does-not-exist";

/// Handshake facts read straight out of a raw `initialize` result, before the
/// list follow-ups. Split from the live capture so the field extraction is
/// unit-tested without a socket.
#[derive(Debug, Clone, PartialEq, Eq, Default)]
struct HandshakeFacts {
    /// `protocolVersion` from the result, if present.
    negotiated_version: Option<String>,
    /// Whether the result carried a `capabilities` object.
    capabilities_advertised: bool,
    /// Whether `capabilities.tools` was advertised, so a `tools/list` MUST
    /// applies.
    advertises_tools: bool,
    /// Whether `capabilities.resources` was advertised.
    advertises_resources: bool,
    /// Whether `capabilities.prompts` was advertised.
    advertises_prompts: bool,
}

/// Turn a `--server` string into a [`ServerSpec`].
///
/// An `http://` or `https://` value is a streamable-HTTP URL; anything else is
/// a stdio command line split on whitespace (`command arg1 arg2`). Those are
/// the two transports a single developer reaches for when probing their own
/// server.
pub(crate) fn parse_server_spec(server: &str) -> ServerSpec {
    let trimmed = server.trim();
    if trimmed.starts_with("http://") || trimmed.starts_with("https://") {
        return ServerSpec::StreamableHttp {
            url: trimmed.to_string(),
            headers: BTreeMap::new(),
            oauth: None,
            stateless: false,
        };
    }
    let mut parts = trimmed.split_whitespace().map(str::to_string);
    let command = parts.next().unwrap_or_default();
    ServerSpec::Stdio {
        command,
        args: parts.collect(),
        env: BTreeMap::new(),
    }
}

/// Pull the probe-relevant facts out of a raw `initialize` result.
fn handshake_facts(raw: &Value) -> HandshakeFacts {
    let negotiated_version = raw
        .get("protocolVersion")
        .and_then(Value::as_str)
        .map(str::to_string);
    let capabilities = raw.get("capabilities");
    let advertises = |name: &str| capabilities.and_then(|c| c.get(name)).is_some();
    HandshakeFacts {
        negotiated_version,
        capabilities_advertised: capabilities.map(Value::is_object).unwrap_or(false),
        advertises_tools: advertises("tools"),
        advertises_resources: advertises("resources"),
        advertises_prompts: advertises("prompts"),
    }
}

/// Issue a `*/list` request and report whether it returned without a JSON-RPC
/// error.
async fn list_succeeds(client: &Client, method: &str) -> bool {
    client.request(method, json!({})).await.is_ok()
}

/// Read a deliberately non-existent resource and record what the server did,
/// the input to the SEP-2164 probes.
async fn read_bogus_resource(client: &Client) -> ResourceReadObservation {
    match client
        .request("resources/read", json!({ "uri": BOGUS_RESOURCE_URI }))
        .await
    {
        Ok(result) => {
            let empty_contents = result
                .get("contents")
                .and_then(Value::as_array)
                .map(|c| c.is_empty())
                .unwrap_or(false);
            ResourceReadObservation {
                is_error: false,
                error_code: None,
                empty_contents,
            }
        }
        Err(ProtocolError::ServerError { code, .. }) => ResourceReadObservation {
            is_error: true,
            error_code: Some(code),
            empty_contents: false,
        },
        // A transport-level failure is not a JSON-RPC error verdict; record it
        // as errored-without-a-code so the MUST (no empty contents) still holds
        // and the SHOULD (-32602) is reported unmet.
        Err(_) => ResourceReadObservation {
            is_error: true,
            error_code: None,
            empty_contents: false,
        },
    }
}

/// Whether the connection still answers a request. A clean JSON-RPC error (an
/// unknown method, say) still counts as responsive; only a transport death or
/// timeout means the server stopped answering.
async fn still_responsive(client: &Client) -> bool {
    matches!(
        client.request("ping", json!({})).await,
        Ok(_) | Err(ProtocolError::ServerError { .. })
    )
}

/// Probe how the server reacts to stray `notifications/cancelled` messages: it
/// must ignore a cancellation for a request id it has no in-flight work for and
/// keep answering, rather than hang or drop the connection.
async fn capture_cancellation(client: &Client) -> CancellationObservation {
    // A request id the client never issued.
    let _ = client
        .notify(
            "notifications/cancelled",
            json!({ "requestId": "mcptest-cancellation-probe-unknown", "reason": "conformance probe" }),
        )
        .await;
    let healthy_after_unknown_cancel = still_responsive(client).await;

    // Cancel a request that already completed. `request_with_id` only yields the
    // id on a clean result, so a server without `ping` leaves this not-probed.
    let completed_cancel = match client.request_with_id("ping", json!({})).await {
        Ok((id, _)) => {
            let _ = client
                .notify(
                    "notifications/cancelled",
                    json!({ "requestId": id, "reason": "conformance probe" }),
                )
                .await;
            Some(still_responsive(client).await)
        }
        Err(_) => None,
    };

    CancellationObservation {
        healthy_after_unknown_cancel,
        completed_cancel,
    }
}

/// Capture handshake evidence from a live server.
///
/// A transport-level connect failure (the server is unreachable) is an `Err`,
/// since there is nothing to score. A handshake that completes but is malformed
/// shows up in the evidence fields and is scored by the probe; a handshake that
/// errors mid-flight is recorded as failed-handshake evidence rather than an
/// error, because "the server cannot complete `initialize`" is itself a real
/// conformance result.
pub(crate) async fn capture_evidence(spec: &ServerSpec) -> Result<ConformanceEvidence> {
    let server = match connect_server(
        "conformance-target",
        spec,
        "mcptest",
        mcptest_core::version::semver(),
        &ProxyConfig::default(),
        &[],
    )
    .await
    {
        Ok(s) => s,
        Err(ConnectError::Initialize(_)) => return Ok(ConformanceEvidence::default()),
        Err(other) => return Err(other).context("connecting to the server under test"),
    };

    let evidence = evidence_from_client(&server.client).await;
    server.shutdown.shutdown().await.ok();
    Ok(evidence)
}

/// Build conformance evidence from a connected client whose handshake already
/// completed.
///
/// `pub(crate)` so the run's conformance lane
/// (`cli::handlers::run_live_lane_conformance`, WOR-1484) reuses this exact
/// probe sequence against a server it connected through the suite's
/// [`crate::commands::compliance_suite::connect_suite_servers`] pool, instead of
/// reconnecting its own. The caller owns the connection lifecycle (shutdown);
/// this only reads. A connected client means `initialize` completed, so
/// `initialize_ok` is always true here. When the raw initialize result was not
/// retained, only that one fact is recorded.
pub(crate) async fn evidence_from_client(client: &Client) -> ConformanceEvidence {
    let Some(result) = client.initialize_response().await else {
        // The handshake completed (the client is connected), but the raw
        // result was not retained. Record the one fact we know.
        return ConformanceEvidence {
            initialize_ok: true,
            ..ConformanceEvidence::default()
        };
    };
    let facts = handshake_facts(&result);
    let tools_list_ok = match facts.advertises_tools {
        true => Some(list_succeeds(client, "tools/list").await),
        false => None,
    };
    let resources_list_ok = match facts.advertises_resources {
        true => Some(list_succeeds(client, "resources/list").await),
        false => None,
    };
    let prompts_list_ok = match facts.advertises_prompts {
        true => Some(list_succeeds(client, "prompts/list").await),
        false => None,
    };
    // SEP-2164 only applies to a server that serves resources.
    let bogus_resource_read = match facts.advertises_resources {
        true => Some(read_bogus_resource(client).await),
        false => None,
    };
    // Cancellation resilience applies to every server (stable since
    // 2024-11-05), so it is probed unconditionally once connected.
    let cancellation = Some(capture_cancellation(client).await);
    ConformanceEvidence {
        initialize_ok: true,
        negotiated_version: facts.negotiated_version,
        capabilities_advertised: facts.capabilities_advertised,
        tools_list_ok,
        resources_list_ok,
        prompts_list_ok,
        bogus_resource_read,
        cancellation,
    }
}

/// Synchronous entry point: build a current-thread runtime and capture
/// evidence from `server`.
pub(crate) fn probe_server_evidence(server: &str) -> Result<ConformanceEvidence> {
    let spec = parse_server_spec(server);
    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .context("failed to build tokio runtime")?;
    runtime.block_on(capture_evidence(&spec))
}

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

    #[test]
    fn http_url_parses_as_streamable_http() {
        let spec = parse_server_spec("https://example.test/mcp");
        match spec {
            ServerSpec::StreamableHttp { url, .. } => {
                assert_eq!(url, "https://example.test/mcp")
            }
            other => panic!("expected streamable-http, got {other:?}"),
        }
    }

    #[test]
    fn plain_http_url_parses_as_streamable_http() {
        let spec = parse_server_spec("http://127.0.0.1:8080/mcp");
        assert!(matches!(spec, ServerSpec::StreamableHttp { .. }));
    }

    #[test]
    fn a_command_line_parses_as_stdio_with_args() {
        let spec = parse_server_spec("python server.py --flag");
        match spec {
            ServerSpec::Stdio { command, args, .. } => {
                assert_eq!(command, "python");
                assert_eq!(args, vec!["server.py", "--flag"]);
            }
            other => panic!("expected stdio, got {other:?}"),
        }
    }

    #[test]
    fn handshake_facts_reads_version_and_capabilities() {
        let raw = json!({
            "protocolVersion": "2025-06-18",
            "capabilities": { "tools": {}, "resources": {} }
        });
        let facts = handshake_facts(&raw);
        assert_eq!(facts.negotiated_version.as_deref(), Some("2025-06-18"));
        assert!(facts.capabilities_advertised);
        assert!(facts.advertises_tools);
    }

    #[test]
    fn handshake_facts_without_tools_capability_does_not_require_tools_list() {
        let raw = json!({
            "protocolVersion": "2025-06-18",
            "capabilities": { "resources": {} }
        });
        let facts = handshake_facts(&raw);
        assert!(facts.capabilities_advertised);
        assert!(!facts.advertises_tools);
    }

    #[test]
    fn handshake_facts_with_no_capabilities_object_is_not_advertised() {
        let raw = json!({ "protocolVersion": "2025-06-18" });
        let facts = handshake_facts(&raw);
        assert!(!facts.capabilities_advertised);
        assert!(!facts.advertises_tools);
        assert_eq!(facts.negotiated_version.as_deref(), Some("2025-06-18"));
    }
}

#[cfg(all(test, unix))]
mod live_tests {
    use super::*;

    use std::os::unix::fs::PermissionsExt;
    use std::path::{Path, PathBuf};

    use mcptest_core::conformance::{must_tally, score_tier, universal_probes, Tier, TierInput};
    use tempfile::TempDir;

    fn write_script(dir: &TempDir, name: &str, body: &str) -> PathBuf {
        let path = dir.path().join(name);
        std::fs::write(&path, body).expect("write script");
        let mut perms = std::fs::metadata(&path).expect("meta").permissions();
        perms.set_mode(0o755);
        std::fs::set_permissions(&path, perms).expect("chmod");
        path
    }

    /// Stdio MCP fixture that advertises the `tools` capability and serves a
    /// one-tool `tools/list`. A clean server: every universal MUST should pass.
    fn write_clean_server(dir: &TempDir) -> PathBuf {
        let body = r##"#!/bin/sh
while IFS= read -r line; do
  case "$line" in
    *'"id":'*) tmp=${line#*'"id":'}; id=${tmp%%[!0-9]*} ;;
    *) id= ;;
  esac
  case "$line" in
    *'"method":"initialize"'*)
      printf '{"jsonrpc":"2.0","id":%s,"result":{"protocolVersion":"2025-03-26","serverInfo":{"name":"clean","version":"0.0.1"},"capabilities":{"tools":{}}}}\n' "$id" ;;
    *'"method":"notifications/initialized"'*) : ;;
    *'"method":"tools/list"'*)
      printf '{"jsonrpc":"2.0","id":%s,"result":{"tools":[{"name":"ping","inputSchema":{"type":"object"}}]}}\n' "$id" ;;
    *)
      [ -n "$id" ] && printf '{"jsonrpc":"2.0","id":%s,"error":{"code":-32601,"message":"method not found"}}\n' "$id" ;;
  esac
done
"##;
        write_script(dir, "clean.sh", body)
    }

    /// Stdio fixture that advertises an empty capability set, so the
    /// `tools/list` MUST does not apply and is omitted from the tally.
    fn write_no_tools_server(dir: &TempDir) -> PathBuf {
        let body = r##"#!/bin/sh
while IFS= read -r line; do
  case "$line" in
    *'"id":'*) tmp=${line#*'"id":'}; id=${tmp%%[!0-9]*} ;;
    *) id= ;;
  esac
  case "$line" in
    *'"method":"initialize"'*)
      printf '{"jsonrpc":"2.0","id":%s,"result":{"protocolVersion":"2025-03-26","serverInfo":{"name":"bare","version":"0.0.1"},"capabilities":{}}}\n' "$id" ;;
    *'"method":"notifications/initialized"'*) : ;;
    *)
      [ -n "$id" ] && printf '{"jsonrpc":"2.0","id":%s,"error":{"code":-32601,"message":"method not found"}}\n' "$id" ;;
  esac
done
"##;
        write_script(dir, "bare.sh", body)
    }

    fn stdio_spec(path: &Path) -> ServerSpec {
        ServerSpec::Stdio {
            command: path.to_string_lossy().into_owned(),
            args: Vec::new(),
            env: BTreeMap::new(),
        }
    }

    #[tokio::test]
    async fn a_clean_stdio_server_probes_to_tier_one() {
        let tmp = TempDir::new().unwrap();
        let spec = stdio_spec(&write_clean_server(&tmp));
        let evidence = capture_evidence(&spec).await.expect("probe a clean server");

        assert!(evidence.initialize_ok);
        assert!(evidence.capabilities_advertised);
        assert_eq!(evidence.negotiated_version.as_deref(), Some("2025-03-26"));
        assert_eq!(evidence.tools_list_ok, Some(true));

        let outcomes = universal_probes("", &evidence);
        assert_eq!(must_tally(&outcomes), (4, 4));
        let tier = score_tier(TierInput {
            must_passed: 4,
            must_total: 4,
            should_passed: 0,
            should_total: 0,
        });
        assert_eq!(tier, Tier::Tier1);
    }

    #[tokio::test]
    async fn a_server_without_tools_omits_the_tools_list_must() {
        let tmp = TempDir::new().unwrap();
        let spec = stdio_spec(&write_no_tools_server(&tmp));
        let evidence = capture_evidence(&spec).await.expect("probe a bare server");

        assert!(evidence.initialize_ok);
        assert!(evidence.capabilities_advertised);
        assert_eq!(evidence.tools_list_ok, None);

        let outcomes = universal_probes("", &evidence);
        assert_eq!(must_tally(&outcomes), (3, 3));
    }

    /// Stdio fixture advertising tools + resources + prompts that serves all
    /// three lists and returns -32602 for any `resources/read`. Exercises the
    /// full per-SEP capture path (WOR-1404).
    fn write_full_server(dir: &TempDir) -> PathBuf {
        let body = r##"#!/bin/sh
while IFS= read -r line; do
  case "$line" in
    *'"id":'*) tmp=${line#*'"id":'}; id=${tmp%%[!0-9]*} ;;
    *) id= ;;
  esac
  case "$line" in
    *'"method":"initialize"'*)
      printf '{"jsonrpc":"2.0","id":%s,"result":{"protocolVersion":"2025-03-26","serverInfo":{"name":"full","version":"0.0.1"},"capabilities":{"tools":{},"resources":{},"prompts":{}}}}\n' "$id" ;;
    *'"method":"notifications/initialized"'*) : ;;
    *'"method":"tools/list"'*)
      printf '{"jsonrpc":"2.0","id":%s,"result":{"tools":[{"name":"ping","inputSchema":{"type":"object"}}]}}\n' "$id" ;;
    *'"method":"resources/list"'*)
      printf '{"jsonrpc":"2.0","id":%s,"result":{"resources":[]}}\n' "$id" ;;
    *'"method":"prompts/list"'*)
      printf '{"jsonrpc":"2.0","id":%s,"result":{"prompts":[]}}\n' "$id" ;;
    *'"method":"resources/read"'*)
      printf '{"jsonrpc":"2.0","id":%s,"error":{"code":-32602,"message":"resource not found"}}\n' "$id" ;;
    *)
      [ -n "$id" ] && printf '{"jsonrpc":"2.0","id":%s,"error":{"code":-32601,"message":"method not found"}}\n' "$id" ;;
  esac
done
"##;
        write_script(dir, "full.sh", body)
    }

    #[tokio::test]
    async fn a_full_server_passes_capability_and_sep_2164_probes() {
        use mcptest_core::conformance::{extended_probes, tally_by_level};
        let tmp = TempDir::new().unwrap();
        let spec = stdio_spec(&write_full_server(&tmp));
        let evidence = capture_evidence(&spec).await.expect("probe a full server");

        assert_eq!(evidence.resources_list_ok, Some(true));
        assert_eq!(evidence.prompts_list_ok, Some(true));
        let obs = evidence
            .bogus_resource_read
            .as_ref()
            .expect("resources advertised, so a bogus read was probed");
        assert!(obs.is_error);
        assert_eq!(obs.error_code, Some(-32602));

        let mut outcomes = universal_probes("", &evidence);
        outcomes.extend(extended_probes(&evidence));
        // 4 universal + 2 capability + 1 SEP-2164 MUST + 1 cancellation MUST
        // (unknown id; the fixture has no `ping`, so the completed-id check is
        // not probed); 1 SEP-2164 SHOULD.
        let (mp, mt, sp, st) = tally_by_level(&outcomes);
        assert_eq!((mp, mt), (8, 8));
        assert_eq!((sp, st), (1, 1));
    }

    /// Stdio fixture that answers `ping` and ignores notifications (so a stray
    /// `notifications/cancelled` is dropped and the connection stays healthy).
    fn write_pinger(dir: &TempDir) -> PathBuf {
        let body = r##"#!/bin/sh
while IFS= read -r line; do
  case "$line" in
    *'"id":'*) tmp=${line#*'"id":'}; id=${tmp%%[!0-9]*} ;;
    *) id= ;;
  esac
  case "$line" in
    *'"method":"initialize"'*)
      printf '{"jsonrpc":"2.0","id":%s,"result":{"protocolVersion":"2025-03-26","serverInfo":{"name":"pinger","version":"0.0.1"},"capabilities":{"tools":{}}}}\n' "$id" ;;
    *'"method":"notifications/'*) : ;;
    *'"method":"ping"'*)
      printf '{"jsonrpc":"2.0","id":%s,"result":{}}\n' "$id" ;;
    *'"method":"tools/list"'*)
      printf '{"jsonrpc":"2.0","id":%s,"result":{"tools":[]}}\n' "$id" ;;
    *)
      [ -n "$id" ] && printf '{"jsonrpc":"2.0","id":%s,"error":{"code":-32601,"message":"method not found"}}\n' "$id" ;;
  esac
done
"##;
        write_script(dir, "pinger.sh", body)
    }

    #[tokio::test]
    async fn a_server_that_ignores_stray_cancellations_is_healthy() {
        use mcptest_core::conformance::cancellation_probes;
        let tmp = TempDir::new().unwrap();
        let spec = stdio_spec(&write_pinger(&tmp));
        let evidence = capture_evidence(&spec).await.expect("probe the pinger");

        let c = evidence.cancellation.as_ref().expect("cancellation probed");
        assert!(c.healthy_after_unknown_cancel);
        // The fixture answers ping, so the completed-id check runs and passes.
        assert_eq!(c.completed_cancel, Some(true));

        let outcomes = cancellation_probes(&evidence);
        assert_eq!(outcomes.len(), 2);
        assert!(outcomes.iter().all(|o| o.result.passed()));
    }

    #[tokio::test]
    async fn an_unreachable_server_is_an_error_not_a_silent_pass() {
        // A command that exits immediately cannot complete the handshake.
        let spec = ServerSpec::Stdio {
            command: "false".to_string(),
            args: Vec::new(),
            env: BTreeMap::new(),
        };
        let result = capture_evidence(&spec).await;
        match result {
            // Either the transport never came up (Err) or the handshake failed
            // and was recorded as not-ok evidence. Both are honest; neither is a
            // silent pass.
            Err(_) => {}
            Ok(evidence) => assert!(
                !evidence.initialize_ok,
                "a server that cannot handshake must not look initialized"
            ),
        }
    }
}
