//! Integration tests for the exec IPC dispatcher.
//!
//! These exercise the full session loop end to end: a duplex pipe
//! stands in for the SDK side, an in-memory backend stands in for the
//! upstream MCP server. The shape of these tests doubles as a
//! reference for how a real SDK speaks to the co-process.

use std::sync::Mutex;

use async_trait::async_trait;
use mcptest_core::exec::{
    error_codes, run, Dispatcher, IpcError, IpcRequest, IpcResponse, IpcResponseBody, McpBackend,
    SessionOptions, COPROCESS_PROTOCOL_VERSION, IPC_VERSION,
};
use serde_json::{json, Value};
use tokio::io::{duplex, AsyncReadExt, AsyncWriteExt, BufReader};

#[derive(Default)]
struct CountingBackend {
    calls: Mutex<Vec<String>>,
    fail_call: Mutex<bool>,
    disconnected: Mutex<bool>,
}

impl CountingBackend {
    fn record(&self, name: &str) {
        self.calls.lock().unwrap().push(name.into());
    }
}

#[async_trait]
impl McpBackend for CountingBackend {
    async fn initialize(&self) -> Result<Value, mcptest_core::exec::dispatcher::BackendError> {
        self.record("initialize");
        Ok(json!({"name": "test-server", "version": "1.2.3"}))
    }
    async fn call_tool(
        &self,
        tool: &str,
        arguments: Value,
    ) -> Result<Value, mcptest_core::exec::dispatcher::BackendError> {
        self.record(&format!("call_tool:{tool}"));
        if *self.disconnected.lock().unwrap() {
            return Err(mcptest_core::exec::dispatcher::BackendError::Disconnected {
                reason: "server pipe closed".into(),
            });
        }
        if *self.fail_call.lock().unwrap() {
            return Err(mcptest_core::exec::dispatcher::BackendError::Server {
                message: "tool refused".into(),
                data: Some(json!({"tool": tool})),
            });
        }
        Ok(json!({
            "content": [{"type": "text", "text": format!("ran {tool}")}],
            "isError": false,
            "_args": arguments,
        }))
    }
    async fn list_tools(&self) -> Result<Value, mcptest_core::exec::dispatcher::BackendError> {
        self.record("list_tools");
        Ok(json!({"tools": [{"name": "echo"}]}))
    }
    async fn list_resources(&self) -> Result<Value, mcptest_core::exec::dispatcher::BackendError> {
        self.record("list_resources");
        Ok(json!({"resources": []}))
    }
    async fn read_resource(
        &self,
        uri: &str,
    ) -> Result<Value, mcptest_core::exec::dispatcher::BackendError> {
        self.record(&format!("read_resource:{uri}"));
        Ok(json!({"contents": [{"uri": uri, "text": "hello"}]}))
    }
    async fn call_prompt(
        &self,
        name: &str,
        _arguments: Value,
    ) -> Result<Value, mcptest_core::exec::dispatcher::BackendError> {
        self.record(&format!("call_prompt:{name}"));
        Ok(json!({"messages": [{"role": "user", "content": "hi"}]}))
    }
    async fn reporter_set_upload(
        &self,
        _url: &str,
        _token: Option<&str>,
    ) -> Result<(), mcptest_core::exec::dispatcher::BackendError> {
        self.record("reporter_set_upload");
        Ok(())
    }
    async fn reporter_flush(&self) -> Result<(), mcptest_core::exec::dispatcher::BackendError> {
        self.record("reporter_flush");
        Ok(())
    }
    async fn cache_set_mode(
        &self,
        mode: &str,
    ) -> Result<(), mcptest_core::exec::dispatcher::BackendError> {
        self.record(&format!("cache:{mode}"));
        Ok(())
    }
    async fn cassette_set_mode(
        &self,
        mode: &str,
    ) -> Result<(), mcptest_core::exec::dispatcher::BackendError> {
        self.record(&format!("cassette:{mode}"));
        Ok(())
    }
    async fn shutdown(&self) -> Result<(), mcptest_core::exec::dispatcher::BackendError> {
        self.record("shutdown");
        Ok(())
    }
}

async fn drain_lines<R: tokio::io::AsyncRead + Unpin>(reader: &mut R, count: usize) -> Vec<String> {
    let mut buf = Vec::new();
    let mut out = Vec::new();
    let mut tmp = [0u8; 4096];
    loop {
        if out.len() >= count {
            break;
        }
        let n = reader.read(&mut tmp).await.unwrap();
        if n == 0 {
            break;
        }
        buf.extend_from_slice(&tmp[..n]);
        while let Some(idx) = buf.iter().position(|b| *b == b'\n') {
            let line = String::from_utf8(buf[..idx].to_vec()).unwrap();
            buf.drain(..=idx);
            out.push(line);
            if out.len() >= count {
                break;
            }
        }
    }
    out
}

fn parse_response(line: &str) -> IpcResponse {
    serde_json::from_str(line).expect("response must parse")
}

async fn run_session(
    backend: CountingBackend,
) -> (
    tokio::io::DuplexStream,
    tokio::io::DuplexStream,
    tokio::task::JoinHandle<std::io::Result<()>>,
) {
    let (sdk_writer, co_reader) = duplex(64 * 1024);
    let (co_writer, sdk_reader) = duplex(64 * 1024);
    let dispatcher = Dispatcher::new(backend);
    let handle = tokio::spawn(async move {
        run(
            BufReader::new(co_reader),
            co_writer,
            dispatcher,
            SessionOptions::default(),
        )
        .await
    });
    (sdk_writer, sdk_reader, handle)
}

#[tokio::test]
async fn initialize_handshake_round_trips() {
    let (mut sdk_w, mut sdk_r, handle) = run_session(CountingBackend::default()).await;
    let req = IpcRequest {
        jsonrpc: "2.0".into(),
        id: Some(json!(1)),
        method: "mcp.initialize".into(),
        params: None,
    };
    sdk_w
        .write_all(format!("{}\n", serde_json::to_string(&req).unwrap()).as_bytes())
        .await
        .unwrap();
    sdk_w
        .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"mcp.shutdown\"}\n")
        .await
        .unwrap();
    let lines = drain_lines(&mut sdk_r, 2).await;
    let init = parse_response(&lines[0]);
    match init.body {
        IpcResponseBody::Success { result } => {
            assert_eq!(result["ipc_version"], IPC_VERSION);
            assert_eq!(result["server"]["name"], "test-server");
        }
        IpcResponseBody::Failure { error } => panic!("init failed: {error:?}"),
    }
    handle.await.unwrap().unwrap();
}

#[tokio::test]
async fn mcp_call_happy_path_returns_tool_result() {
    let (mut sdk_w, mut sdk_r, handle) = run_session(CountingBackend::default()).await;
    let req = IpcRequest {
        jsonrpc: "2.0".into(),
        id: Some(json!(11)),
        method: "mcp.call".into(),
        params: Some(json!({"tool": "echo", "arguments": {"hello": "world"}})),
    };
    sdk_w
        .write_all(format!("{}\n", serde_json::to_string(&req).unwrap()).as_bytes())
        .await
        .unwrap();
    sdk_w
        .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":12,\"method\":\"mcp.shutdown\"}\n")
        .await
        .unwrap();
    let lines = drain_lines(&mut sdk_r, 2).await;
    let resp = parse_response(&lines[0]);
    assert_eq!(resp.id, json!(11));
    match resp.body {
        IpcResponseBody::Success { result } => {
            assert_eq!(result["content"][0]["text"], "ran echo");
            assert_eq!(result["_args"]["hello"], "world");
        }
        IpcResponseBody::Failure { error } => panic!("call failed: {error:?}"),
    }
    handle.await.unwrap().unwrap();
}

#[tokio::test]
async fn mcp_call_with_server_error_maps_to_mcp_server_error_code() {
    let backend = CountingBackend::default();
    *backend.fail_call.lock().unwrap() = true;
    let (mut sdk_w, mut sdk_r, handle) = run_session(backend).await;
    sdk_w
        .write_all(
            br#"{"jsonrpc":"2.0","id":1,"method":"mcp.call","params":{"tool":"echo","arguments":{}}}
"#,
        )
        .await
        .unwrap();
    sdk_w
        .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"mcp.shutdown\"}\n")
        .await
        .unwrap();
    let lines = drain_lines(&mut sdk_r, 2).await;
    let resp = parse_response(&lines[0]);
    match resp.body {
        IpcResponseBody::Failure { error } => {
            assert_eq!(error.code, error_codes::MCP_SERVER_ERROR);
            assert_eq!(error.message, "tool refused");
        }
        IpcResponseBody::Success { result } => panic!("expected failure, got {result:?}"),
    }
    handle.await.unwrap().unwrap();
}

#[tokio::test]
async fn server_disconnect_mid_call_surfaces_disconnected_code() {
    let backend = CountingBackend::default();
    *backend.disconnected.lock().unwrap() = true;
    let (mut sdk_w, mut sdk_r, handle) = run_session(backend).await;
    sdk_w
        .write_all(
            br#"{"jsonrpc":"2.0","id":7,"method":"mcp.call","params":{"tool":"echo"}}
"#,
        )
        .await
        .unwrap();
    sdk_w
        .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":8,\"method\":\"mcp.shutdown\"}\n")
        .await
        .unwrap();
    let lines = drain_lines(&mut sdk_r, 2).await;
    let resp = parse_response(&lines[0]);
    match resp.body {
        IpcResponseBody::Failure { error } => {
            assert_eq!(error.code, error_codes::MCP_DISCONNECTED);
        }
        IpcResponseBody::Success { result } => panic!("expected failure, got {result:?}"),
    }
    handle.await.unwrap().unwrap();
}

#[tokio::test]
async fn malformed_ipc_returns_parse_error_then_keeps_serving() {
    let (mut sdk_w, mut sdk_r, handle) = run_session(CountingBackend::default()).await;
    sdk_w.write_all(b"this-is-not-json\n").await.unwrap();
    sdk_w
        .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"mcp.listTools\"}\n")
        .await
        .unwrap();
    sdk_w
        .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"mcp.shutdown\"}\n")
        .await
        .unwrap();
    let lines = drain_lines(&mut sdk_r, 3).await;
    let bad = parse_response(&lines[0]);
    match bad.body {
        IpcResponseBody::Failure { error } => {
            assert_eq!(error.code, error_codes::PARSE_ERROR);
        }
        IpcResponseBody::Success { result } => panic!("expected parse error, got {result:?}"),
    }
    let listed = parse_response(&lines[1]);
    match listed.body {
        IpcResponseBody::Success { result } => {
            assert!(result["tools"].is_array());
        }
        IpcResponseBody::Failure { error } => panic!("list_tools failed: {error:?}"),
    }
    handle.await.unwrap().unwrap();
}

#[tokio::test]
async fn shutdown_flushes_reporter_and_exits() {
    let (mut sdk_w, mut sdk_r, handle) = run_session(CountingBackend::default()).await;
    sdk_w
        .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"mcp.reporter.flush\"}\n")
        .await
        .unwrap();
    sdk_w
        .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"mcp.shutdown\"}\n")
        .await
        .unwrap();
    let lines = drain_lines(&mut sdk_r, 2).await;
    let flushed = parse_response(&lines[0]);
    match flushed.body {
        IpcResponseBody::Success { result } => assert_eq!(result["ok"], true),
        IpcResponseBody::Failure { error } => panic!("flush failed: {error:?}"),
    }
    let bye = parse_response(&lines[1]);
    match bye.body {
        IpcResponseBody::Success { result } => assert_eq!(result["ok"], true),
        IpcResponseBody::Failure { error } => panic!("shutdown failed: {error:?}"),
    }
    handle.await.unwrap().unwrap();
}

#[tokio::test]
async fn unknown_method_returns_method_not_found() {
    let (mut sdk_w, mut sdk_r, handle) = run_session(CountingBackend::default()).await;
    sdk_w
        .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"mcp.nope\"}\n")
        .await
        .unwrap();
    sdk_w
        .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"mcp.shutdown\"}\n")
        .await
        .unwrap();
    let lines = drain_lines(&mut sdk_r, 2).await;
    let resp = parse_response(&lines[0]);
    match resp.body {
        IpcResponseBody::Failure { error } => {
            assert_eq!(error.code, error_codes::METHOD_NOT_FOUND);
        }
        IpcResponseBody::Success { result } => panic!("expected failure, got {result:?}"),
    }
    handle.await.unwrap().unwrap();
}

#[tokio::test]
async fn stdin_eof_triggers_clean_shutdown() {
    let (sdk_w, mut sdk_r, handle) = run_session(CountingBackend::default()).await;
    drop(sdk_w);
    let mut buf = [0u8; 16];
    let n = sdk_r.read(&mut buf).await.unwrap_or(0);
    assert_eq!(n, 0);
    handle.await.unwrap().unwrap();
}

#[tokio::test]
async fn cache_and_cassette_modes_round_trip() {
    let (mut sdk_w, mut sdk_r, handle) = run_session(CountingBackend::default()).await;
    sdk_w
        .write_all(
            br#"{"jsonrpc":"2.0","id":1,"method":"mcp.cache.set_mode","params":{"mode":"off"}}
"#,
        )
        .await
        .unwrap();
    sdk_w
        .write_all(
            br#"{"jsonrpc":"2.0","id":2,"method":"mcp.cassette.set_mode","params":{"mode":"replay"}}
"#,
        )
        .await
        .unwrap();
    sdk_w
        .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"mcp.shutdown\"}\n")
        .await
        .unwrap();
    let lines = drain_lines(&mut sdk_r, 3).await;
    let cache = parse_response(&lines[0]);
    let cassette = parse_response(&lines[1]);
    match cache.body {
        IpcResponseBody::Success { result } => assert_eq!(result["mode"], "off"),
        IpcResponseBody::Failure { error } => panic!("cache failed: {error:?}"),
    }
    match cassette.body {
        IpcResponseBody::Success { result } => assert_eq!(result["mode"], "replay"),
        IpcResponseBody::Failure { error } => panic!("cassette failed: {error:?}"),
    }
    handle.await.unwrap().unwrap();
}

#[tokio::test]
async fn invalid_jsonrpc_version_returns_invalid_request() {
    let (mut sdk_w, mut sdk_r, handle) = run_session(CountingBackend::default()).await;
    sdk_w
        .write_all(b"{\"jsonrpc\":\"1.0\",\"id\":1,\"method\":\"mcp.initialize\"}\n")
        .await
        .unwrap();
    sdk_w
        .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"mcp.shutdown\"}\n")
        .await
        .unwrap();
    let lines = drain_lines(&mut sdk_r, 2).await;
    let resp = parse_response(&lines[0]);
    match resp.body {
        IpcResponseBody::Failure { error } => {
            assert_eq!(error.code, error_codes::INVALID_REQUEST);
        }
        IpcResponseBody::Success { result } => panic!("expected invalid request, got {result:?}"),
    }
    handle.await.unwrap().unwrap();
}

#[tokio::test]
async fn handshake_round_trips_versions_and_binary_semver() {
    let (mut sdk_w, mut sdk_r, handle) = run_session(CountingBackend::default()).await;
    let req = IpcRequest {
        jsonrpc: "2.0".into(),
        id: Some(json!(1)),
        method: "coprocess/handshake".into(),
        params: Some(json!({
            "protocol_version": COPROCESS_PROTOCOL_VERSION,
            "sdk": "rust",
            "sdk_version": "1.0.1",
        })),
    };
    sdk_w
        .write_all(format!("{}\n", serde_json::to_string(&req).unwrap()).as_bytes())
        .await
        .unwrap();
    sdk_w
        .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"mcp.shutdown\"}\n")
        .await
        .unwrap();
    let lines = drain_lines(&mut sdk_r, 2).await;
    let resp = parse_response(&lines[0]);
    match resp.body {
        IpcResponseBody::Success { result } => {
            assert_eq!(result["protocol_version"], COPROCESS_PROTOCOL_VERSION);
            let binary_version = result["binary_version"].as_str().unwrap_or_default();
            assert!(
                !binary_version.is_empty(),
                "handshake must carry the binary semver, got {result}"
            );
        }
        IpcResponseBody::Failure { error } => panic!("handshake failed: {error:?}"),
    }
    handle.await.unwrap().unwrap();
}

#[tokio::test]
async fn handshake_version_mismatch_is_rejected_naming_both_versions() {
    let (mut sdk_w, mut sdk_r, handle) = run_session(CountingBackend::default()).await;
    let newer = COPROCESS_PROTOCOL_VERSION + 1;
    let req = IpcRequest {
        jsonrpc: "2.0".into(),
        id: Some(json!(1)),
        method: "coprocess/handshake".into(),
        params: Some(json!({"protocol_version": newer, "sdk": "python"})),
    };
    sdk_w
        .write_all(format!("{}\n", serde_json::to_string(&req).unwrap()).as_bytes())
        .await
        .unwrap();
    sdk_w
        .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"mcp.shutdown\"}\n")
        .await
        .unwrap();
    let lines = drain_lines(&mut sdk_r, 2).await;
    let resp = parse_response(&lines[0]);
    match resp.body {
        IpcResponseBody::Failure { error } => {
            assert_eq!(error.code, error_codes::INVALID_PARAMS);
            assert!(
                error
                    .message
                    .contains(&format!("speaks v{COPROCESS_PROTOCOL_VERSION}")),
                "got: {}",
                error.message
            );
            assert!(
                error.message.contains(&format!("sent v{newer}")),
                "got: {}",
                error.message
            );
            assert!(
                error.message.contains("Upgrade the mcptest binary"),
                "got: {}",
                error.message
            );
        }
        IpcResponseBody::Success { result } => panic!("expected rejection, got {result:?}"),
    }
    handle.await.unwrap().unwrap();
}

#[tokio::test]
async fn legacy_sdks_that_send_no_handshake_are_still_served() {
    // WOR-1318 enforcement is opt-in by design: a v1 SDK never sends
    // `coprocess/handshake`, so the very first request may be a tool
    // call and the session must serve it.
    let (mut sdk_w, mut sdk_r, handle) = run_session(CountingBackend::default()).await;
    sdk_w
        .write_all(
            br#"{"jsonrpc":"2.0","id":1,"method":"mcp.call","params":{"tool":"echo","arguments":{}}}
"#,
        )
        .await
        .unwrap();
    sdk_w
        .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"mcp.shutdown\"}\n")
        .await
        .unwrap();
    let lines = drain_lines(&mut sdk_r, 2).await;
    let resp = parse_response(&lines[0]);
    match resp.body {
        IpcResponseBody::Success { result } => {
            assert_eq!(result["content"][0]["text"], "ran echo");
        }
        IpcResponseBody::Failure { error } => {
            panic!("handshake-less session must keep working: {error:?}")
        }
    }
    handle.await.unwrap().unwrap();
}

#[tokio::test]
async fn handshake_with_unknown_param_field_is_rejected() {
    // Drift detection: extra handshake params mean a newer SDK contract
    // and must fail loudly rather than parse partially.
    let (mut sdk_w, mut sdk_r, handle) = run_session(CountingBackend::default()).await;
    let params = json!({"protocol_version": COPROCESS_PROTOCOL_VERSION, "transport": "quic"});
    let line = format!(
        "{{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"coprocess/handshake\",\"params\":{params}}}\n"
    );
    sdk_w.write_all(line.as_bytes()).await.unwrap();
    sdk_w
        .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"mcp.shutdown\"}\n")
        .await
        .unwrap();
    let lines = drain_lines(&mut sdk_r, 2).await;
    let resp = parse_response(&lines[0]);
    match resp.body {
        IpcResponseBody::Failure { error } => {
            assert_eq!(error.code, error_codes::INVALID_PARAMS);
        }
        IpcResponseBody::Success { result } => panic!("expected rejection, got {result:?}"),
    }
    handle.await.unwrap().unwrap();
}

/// Cross-check the two `COPROCESS_PROTOCOL_VERSION` literals (WOR-1318).
///
/// The Rust SDK cannot depend on this crate (its only shared workspace
/// dependency is a proc-macro crate, and pulling `mcptest-core` in would
/// drag the async stack into a small sync crate), so the version is
/// mirrored as a second literal and this test is what keeps the mirror
/// honest: it reads the SDK source and fails the gate when the values
/// diverge. The tradeoff is a source-text parse that only runs inside
/// the repo checkout; a published `mcptest-core` tarball has no `sdks/`
/// tree, so the test skips there instead of failing spuriously.
#[test]
fn sdk_and_binary_pin_the_same_coprocess_protocol_version() {
    let sdk_source = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
        .join("../../sdks/rust/mcptest-runner/src/coprocess.rs");
    let Ok(source) = std::fs::read_to_string(&sdk_source) else {
        eprintln!(
            "skipping: SDK source not present at {} (published-crate layout)",
            sdk_source.display()
        );
        return;
    };
    let needle = "pub const COPROCESS_PROTOCOL_VERSION: u32 = ";
    let start = source
        .find(needle)
        .expect("SDK must declare pub const COPROCESS_PROTOCOL_VERSION: u32");
    let rest = &source[start + needle.len()..];
    let end = rest.find(';').expect("const declaration ends with `;`");
    let sdk_version: u32 = rest[..end]
        .trim()
        .parse()
        .expect("SDK version literal parses as u32");
    assert_eq!(
        sdk_version, COPROCESS_PROTOCOL_VERSION,
        "sdks/rust/mcptest-runner and mcptest-core disagree on \
         COPROCESS_PROTOCOL_VERSION; bump both literals together"
    );
}

#[tokio::test]
async fn ipc_error_envelope_preserves_optional_data() {
    let resp = IpcResponse::failure(
        json!("string-id"),
        IpcError {
            code: error_codes::MCP_SERVER_ERROR,
            message: "upstream said no".into(),
            data: Some(json!({"detail": "rate limited"})),
        },
    );
    let raw = serde_json::to_string(&resp).unwrap();
    assert!(raw.contains("string-id"));
    assert!(raw.contains("rate limited"));
}
