//! Integration tests for the protocol [`Client`].
//!
//! Exercises the handshake, concurrent in-flight correlation,
//! out-of-order responses, timeouts, notification routing, and
//! server-initiated request handling against an in-memory mock
//! transport.

use std::sync::Arc;
use std::time::Duration;

use mcptest_core::protocol::{
    Client, JsonRpcErrorBody, JsonRpcId, JsonRpcMessage, ProtocolError, ProtocolVersion,
    ProtocolViolationKind, Transport,
};
use serde_json::{json, Value};
use tokio::sync::{mpsc, Mutex};
use tokio::time::timeout;

/// Build a mock transport pair: a Client-side `Transport` plus a
/// `MockPeer` the test can drive.
fn mock_pair() -> (Transport, MockPeer) {
    let (out_tx, out_rx) = mpsc::unbounded_channel(); // Client -> Peer
    let (in_tx, in_rx) = mpsc::unbounded_channel(); // Peer -> Client
    (
        Transport::new(out_tx, in_rx),
        MockPeer {
            client_outbound: Arc::new(Mutex::new(out_rx)),
            client_inbound: in_tx,
        },
    )
}

/// In-memory peer for tests. Receives messages the Client emits and
/// pushes whatever responses the test scripts back.
struct MockPeer {
    client_outbound: Arc<Mutex<mpsc::UnboundedReceiver<JsonRpcMessage>>>,
    client_inbound: mpsc::UnboundedSender<JsonRpcMessage>,
}

impl MockPeer {
    async fn next_outbound(&self) -> JsonRpcMessage {
        let mut rx = self.client_outbound.lock().await;
        rx.recv()
            .await
            .expect("client emitted at least one envelope")
    }

    fn send(&self, msg: JsonRpcMessage) {
        self.client_inbound.send(msg).expect("inbound channel open");
    }
}

#[tokio::test]
async fn handshake_completes_and_returns_negotiated_version() {
    let (transport, peer) = mock_pair();
    let client = Client::builder(transport).build();

    let peer_handle = tokio::spawn(async move {
        let req = peer.next_outbound().await;
        assert_eq!(req.method.as_deref(), Some("initialize"));
        let id = req.id.clone().expect("initialize carries an id");
        peer.send(JsonRpcMessage::response(
            id,
            json!({
                "protocolVersion": "2025-06-18",
                "capabilities": {"tools": {}},
                "serverInfo": {"name": "mock", "version": "0.1"}
            }),
        ));
        let note = peer.next_outbound().await;
        assert_eq!(note.method.as_deref(), Some("initialized"));
        assert!(note.id.is_none(), "initialized must be a notification");
        peer
    });

    let outcome = client
        .initialize(ProtocolVersion::current())
        .await
        .expect("handshake succeeds");
    assert_eq!(outcome.negotiated_version, ProtocolVersion::V2025_06_18);
    assert_eq!(outcome.server.name.as_deref(), Some("mock"));
    let _peer = peer_handle.await.unwrap();
    client.close().await.unwrap();
}

#[tokio::test]
async fn handshake_accepts_all_four_supported_versions() {
    for wire in ["2024-11-05", "2025-03-26", "2025-06-18", "2026-03-26"] {
        let (transport, peer) = mock_pair();
        let client = Client::builder(transport).build();
        let h = tokio::spawn(async move {
            let req = peer.next_outbound().await;
            let id = req.id.clone().unwrap();
            peer.send(JsonRpcMessage::response(
                id,
                json!({"protocolVersion": wire}),
            ));
            let _ = peer.next_outbound().await;
            peer
        });
        let outcome = client.initialize(ProtocolVersion::current()).await.unwrap();
        assert_eq!(outcome.negotiated_version.as_wire(), wire);
        let _peer = h.await.unwrap();
        client.close().await.unwrap();
    }
}

#[tokio::test]
async fn handshake_errors_on_unsupported_protocol_version() {
    let (transport, peer) = mock_pair();
    let client = Client::builder(transport).build();
    let h = tokio::spawn(async move {
        let req = peer.next_outbound().await;
        let id = req.id.clone().unwrap();
        peer.send(JsonRpcMessage::response(
            id,
            json!({"protocolVersion": "1999-01-01"}),
        ));
        peer
    });
    let err = client
        .initialize(ProtocolVersion::current())
        .await
        .expect_err("must reject");
    assert!(matches!(err, ProtocolError::ProtocolVersionMismatch(v) if v == "1999-01-01"));
    let _peer = h.await.unwrap();
}

#[tokio::test]
async fn concurrent_in_flight_requests_correlate_out_of_order() {
    let (transport, peer) = mock_pair();
    let client = Client::builder(transport).build();

    let peer_handle = tokio::spawn(async move {
        // Drain two requests, then reply to the second one first.
        let first = peer.next_outbound().await;
        let second = peer.next_outbound().await;
        peer.send(JsonRpcMessage::response(
            second.id.clone().unwrap(),
            json!({"who": "second"}),
        ));
        peer.send(JsonRpcMessage::response(
            first.id.clone().unwrap(),
            json!({"who": "first"}),
        ));
        peer
    });

    let c1 = client.clone();
    let c2 = client.clone();
    let req1 = tokio::spawn(async move { c1.request("tools/list", json!({})).await });
    let req2 = tokio::spawn(async move { c2.request("resources/list", json!({})).await });
    let r1 = req1.await.unwrap().unwrap();
    let r2 = req2.await.unwrap().unwrap();
    assert_eq!(r1.get("who").and_then(Value::as_str), Some("first"));
    assert_eq!(r2.get("who").and_then(Value::as_str), Some("second"));
    let _peer = peer_handle.await.unwrap();
    client.close().await.unwrap();
}

#[tokio::test]
async fn request_timeout_fires_when_response_never_arrives() {
    let (transport, peer) = mock_pair();
    let client = Client::builder(transport)
        .with_request_timeout(Duration::from_millis(80))
        .build();
    // Drain the request but never reply, so the Client must time out.
    let _drain = tokio::spawn(async move {
        let _ = peer.next_outbound().await;
        peer
    });
    let err = client
        .request("tools/list", json!({}))
        .await
        .expect_err("must time out");
    assert!(
        matches!(err, ProtocolError::RequestTimeout { timeout_ms, .. } if timeout_ms == 80),
        "got: {err:?}"
    );
    client.close().await.unwrap();
}

#[tokio::test]
async fn server_error_response_surfaces_as_typed_error() {
    let (transport, peer) = mock_pair();
    let client = Client::builder(transport).build();
    let h = tokio::spawn(async move {
        let req = peer.next_outbound().await;
        let id = req.id.clone().unwrap();
        peer.send(JsonRpcMessage::error_response(
            id,
            JsonRpcErrorBody {
                code: -32601,
                message: "method not found".into(),
                data: None,
            },
        ));
        peer
    });
    let err = client
        .request("tools/missing", json!({}))
        .await
        .expect_err("must error");
    assert!(matches!(
        err,
        ProtocolError::ServerError { code: -32601, .. }
    ));
    let _peer = h.await.unwrap();
    client.close().await.unwrap();
}

#[tokio::test]
async fn notifications_are_recorded_when_no_handler_is_registered() {
    let (transport, peer) = mock_pair();
    let client = Client::builder(transport).build();
    peer.send(JsonRpcMessage::notification(
        "notifications/progress",
        json!({"value": 1}),
    ));
    peer.send(JsonRpcMessage::notification(
        "notifications/progress",
        json!({"value": 2}),
    ));
    // Give the pump time to drain both envelopes.
    tokio::time::sleep(Duration::from_millis(40)).await;
    let drained = client.drain_notifications().await;
    assert_eq!(drained.len(), 2);
    assert_eq!(drained[0].0, "notifications/progress");
    assert_eq!(drained[0].1["value"], 1);
    assert_eq!(drained[1].1["value"], 2);
    drop(peer);
    client.close().await.unwrap();
}

#[tokio::test]
async fn notification_handler_receives_every_notification() {
    let (transport, peer) = mock_pair();
    let sink: Arc<Mutex<Vec<(String, Value)>>> = Arc::new(Mutex::new(Vec::new()));
    let sink_clone = Arc::clone(&sink);
    let handler = Arc::new(move |method: String, params: Value| {
        let s = Arc::clone(&sink_clone);
        // Synchronous handler: push into the shared sink.
        tokio::spawn(async move {
            s.lock().await.push((method, params));
        });
    });
    let client = Client::builder(transport)
        .with_notification_handler(handler)
        .build();
    peer.send(JsonRpcMessage::notification("ping", json!({"n": 1})));
    peer.send(JsonRpcMessage::notification("ping", json!({"n": 2})));
    tokio::time::sleep(Duration::from_millis(60)).await;
    let collected = sink.lock().await.clone();
    assert_eq!(collected.len(), 2);
    drop(peer);
    client.close().await.unwrap();
}

#[tokio::test]
async fn server_initiated_request_dispatches_to_user_handler() {
    let (transport, peer) = mock_pair();
    let handler: mcptest_core::protocol::ServerRequestHandler =
        Arc::new(|method: String, params: Value| {
            Box::pin(async move {
                assert_eq!(method, "sampling/createMessage");
                Ok(json!({"echo": params}))
            })
        });
    let client = Client::builder(transport)
        .with_server_request_handler(handler)
        .build();
    peer.send(JsonRpcMessage::request(
        JsonRpcId::Number(42),
        "sampling/createMessage",
        json!({"prompt": "hi"}),
    ));
    let reply = timeout(Duration::from_millis(500), async {
        peer.next_outbound().await
    })
    .await
    .expect("client emitted a reply");
    assert_eq!(reply.id, Some(JsonRpcId::Number(42)));
    let result = reply.result.expect("success response");
    assert_eq!(result["echo"]["prompt"], "hi");
    drop(peer);
    client.close().await.unwrap();
}

#[tokio::test]
async fn server_request_without_handler_returns_method_not_found() {
    let (transport, peer) = mock_pair();
    let client = Client::builder(transport).build();
    peer.send(JsonRpcMessage::request(
        JsonRpcId::Number(7),
        "sampling/createMessage",
        json!({}),
    ));
    let reply = timeout(Duration::from_millis(200), async {
        peer.next_outbound().await
    })
    .await
    .expect("client replied");
    let err = reply.error.expect("error response");
    assert_eq!(err.code, -32601);
    drop(peer);
    client.close().await.unwrap();
}

#[tokio::test]
async fn string_id_echo_resolves_request_and_records_violation() {
    // WOR-1311: a server echoing the numeric id 1 as the string "1"
    // must still resolve the request (instead of timing out) and leave
    // evidence behind for the report.
    let (transport, peer) = mock_pair();
    let client = Client::builder(transport).build();
    let h = tokio::spawn(async move {
        let req = peer.next_outbound().await;
        let numeric = match req.id.clone().expect("request carries an id") {
            JsonRpcId::Number(n) => n,
            other => panic!("client allocates numeric ids, got {other:?}"),
        };
        peer.send(JsonRpcMessage::response(
            JsonRpcId::String(numeric.to_string()),
            json!({"ok": true}),
        ));
        peer
    });
    let result = client
        .request("tools/list", json!({}))
        .await
        .expect("string-id echo must still resolve the request");
    assert_eq!(result["ok"], json!(true));
    let summary = client.protocol_violations().await;
    assert_eq!(summary.total, 1);
    assert_eq!(
        summary.examples[0].kind,
        ProtocolViolationKind::StringIdEcho
    );
    let _peer = h.await.unwrap();
    client.close().await.unwrap();
}

#[tokio::test]
async fn unknown_id_response_records_orphan_violation() {
    // WOR-1311: a response correlating to nothing must not vanish; it
    // is recorded so the run can report the misbehaving server.
    let (transport, peer) = mock_pair();
    let client = Client::builder(transport).build();
    peer.send(JsonRpcMessage::response(JsonRpcId::Number(999), json!({})));
    // Give the pump time to drain the envelope.
    tokio::time::sleep(Duration::from_millis(40)).await;
    let summary = client.protocol_violations().await;
    assert_eq!(summary.total, 1);
    assert_eq!(
        summary.examples[0].kind,
        ProtocolViolationKind::OrphanResponse
    );
    assert!(
        summary.examples[0].detail.contains("999"),
        "detail should name the offending id: {}",
        summary.examples[0].detail
    );
    drop(peer);
    client.close().await.unwrap();
}

#[tokio::test]
async fn malformed_frame_records_violation_with_snippet() {
    // WOR-1311: a frame carrying both `result` and `error` fails
    // envelope classification; the pump must record it with a snippet
    // of the offending frame rather than skip it.
    let (transport, peer) = mock_pair();
    let client = Client::builder(transport).build();
    peer.send(JsonRpcMessage {
        jsonrpc: "2.0".into(),
        id: Some(JsonRpcId::Number(1)),
        method: None,
        params: None,
        result: Some(json!({"ok": true})),
        error: Some(JsonRpcErrorBody {
            code: -32000,
            message: "boom".into(),
            data: None,
        }),
    });
    tokio::time::sleep(Duration::from_millis(40)).await;
    let summary = client.protocol_violations().await;
    assert_eq!(summary.total, 1);
    assert_eq!(
        summary.examples[0].kind,
        ProtocolViolationKind::MalformedFrame
    );
    assert!(
        summary.examples[0].detail.contains("\"jsonrpc\":\"2.0\""),
        "detail should carry a frame snippet: {}",
        summary.examples[0].detail
    );
    drop(peer);
    client.close().await.unwrap();
}

#[tokio::test]
async fn violations_accessor_reports_totals_across_kinds() {
    // WOR-1311: the accessor is the runner's hook for attaching
    // violations to reports, so it must aggregate every kind.
    let (transport, peer) = mock_pair();
    let client = Client::builder(transport).build();
    peer.send(JsonRpcMessage::response(JsonRpcId::Number(500), json!({})));
    peer.send(JsonRpcMessage::response(JsonRpcId::Number(501), json!({})));
    peer.send(JsonRpcMessage {
        jsonrpc: "1.0".into(),
        id: None,
        method: Some("ping".into()),
        params: None,
        result: None,
        error: None,
    });
    tokio::time::sleep(Duration::from_millis(40)).await;
    let summary = client.protocol_violations().await;
    assert_eq!(summary.total, 3);
    let orphans = summary
        .examples
        .iter()
        .filter(|v| v.kind == ProtocolViolationKind::OrphanResponse)
        .count();
    let malformed = summary
        .examples
        .iter()
        .filter(|v| v.kind == ProtocolViolationKind::MalformedFrame)
        .count();
    assert_eq!(orphans, 2);
    assert_eq!(malformed, 1);
    drop(peer);
    client.close().await.unwrap();
}

#[tokio::test]
async fn panicking_server_request_handler_does_not_hang_close() {
    // WOR-1314: handler tasks are tracked in a JoinSet owned by the
    // Client, so a panicking handler is reaped (and logged via its join
    // error) instead of vanishing, and close() still returns promptly.
    let (transport, peer) = mock_pair();
    let handler: mcptest_core::protocol::ServerRequestHandler =
        Arc::new(|_method: String, _params: Value| Box::pin(async move { panic!("handler boom") }));
    let client = Client::builder(transport)
        .with_server_request_handler(handler)
        .build();
    peer.send(JsonRpcMessage::request(
        JsonRpcId::Number(11),
        "sampling/createMessage",
        json!({}),
    ));
    // Give the pump time to spawn the handler and the handler time to
    // panic before close() drains the task set.
    tokio::time::sleep(Duration::from_millis(50)).await;
    timeout(Duration::from_secs(2), client.close())
        .await
        .expect("close must not hang on a panicked handler")
        .expect("close succeeds");
    drop(peer);
}

#[tokio::test]
async fn close_aborts_in_flight_server_request_handler_tasks() {
    // WOR-1314: a handler that never finishes must not keep the
    // process alive past close(); the Client aborts its tracked tasks.
    let (transport, peer) = mock_pair();
    let handler: mcptest_core::protocol::ServerRequestHandler =
        Arc::new(|_method: String, _params: Value| {
            Box::pin(async move {
                tokio::time::sleep(Duration::from_secs(60)).await;
                Ok(json!({}))
            })
        });
    let client = Client::builder(transport)
        .with_server_request_handler(handler)
        .build();
    peer.send(JsonRpcMessage::request(
        JsonRpcId::Number(12),
        "sampling/createMessage",
        json!({}),
    ));
    tokio::time::sleep(Duration::from_millis(50)).await;
    let started = std::time::Instant::now();
    timeout(Duration::from_secs(2), client.close())
        .await
        .expect("close must abort the hanging handler instead of waiting")
        .expect("close succeeds");
    assert!(
        started.elapsed() < Duration::from_secs(2),
        "close should return promptly after aborting handler tasks"
    );
    drop(peer);
}

#[tokio::test]
async fn close_aborts_outstanding_requests() {
    let (transport, peer) = mock_pair();
    let client = Client::builder(transport)
        .with_request_timeout(Duration::from_secs(10))
        .build();
    let c = client.clone();
    let req = tokio::spawn(async move { c.request("tools/list", json!({})).await });
    // Drain the outbound envelope so the request lands in the in-flight map.
    let _drain = peer.next_outbound().await;
    // Now close the client. The in-flight request should abort.
    client.close().await.unwrap();
    let err = req.await.unwrap().expect_err("must abort");
    assert!(matches!(err, ProtocolError::Aborted { .. }));
    drop(peer);
}

#[tokio::test]
async fn send_request_exposes_the_id_then_resolves_the_response() {
    // WOR-1407: the fire-and-observe primitive hands back the request id
    // before the reply arrives, then awaits the correlated response.
    let (transport, peer) = mock_pair();
    let client = Client::builder(transport).build();
    let h = tokio::spawn(async move {
        let req = peer.next_outbound().await;
        assert_eq!(req.method.as_deref(), Some("tools/list"));
        let id = req.id.clone().expect("request carries an id");
        peer.send(JsonRpcMessage::response(id, json!({ "tools": [] })));
        peer
    });

    let pending = client
        .send_request("tools/list", json!({}))
        .await
        .expect("send succeeds");
    assert!(
        pending.id() > 0,
        "the wire id is available before the reply"
    );
    let value = pending.response().await.expect("response resolves");
    assert_eq!(value, json!({ "tools": [] }));
    let _peer = h.await.unwrap();
    client.close().await.unwrap();
}

#[tokio::test]
async fn response_within_returns_none_when_no_reply_arrives() {
    // The cancel-after-completion race probe relies on this: after sending a
    // request the peer never answers, a bounded observe returns None rather
    // than blocking out the full request timeout.
    let (transport, peer) = mock_pair();
    let client = Client::builder(transport)
        .with_request_timeout(Duration::from_secs(30))
        .build();
    let pending = client
        .send_request("tools/call", json!({}))
        .await
        .expect("send succeeds");
    let observed = pending
        .response_within(Duration::from_millis(50))
        .await
        .expect("observe does not error");
    assert!(
        observed.is_none(),
        "no reply within the window must be None"
    );
    drop(peer);
    client.close().await.unwrap();
}
