//! QA-owned protocol corpus.
//!
//! ## Why a second file
//!
//! `tests/protocol_client.rs` already covers the protocol Client's
//! happy-path integrations (handshake, basic correlation, server
//! requests). This file is property-test territory: it sweeps the
//! request/response correlation invariant under randomized permutations,
//! sweeps every supported protocol version through the handshake, and
//! exercises every error envelope variant the Client must surface as a
//! typed `ProtocolError`.
//!
//! ## Scope
//!
//! * Handshake property: initialize/initialized round trip for each of
//!   the four spec revisions.
//! * Concurrent correlation: 10 in-flight requests responded in any
//!   permutation, all routed to the right caller.
//! * Error envelope handling: -32602, -32603, -32700, -32000 each lift
//!   into `ProtocolError::ServerError` with the matching code.
//! * Notification routing: notifications without an id reach the
//!   notification handler if registered and never fan into the
//!   in-flight map.
//! * Capability passthrough: client capabilities round-trip through the
//!   handshake; server capabilities arrive intact on the outcome.
//! * Protocol version mismatch: clear `ProtocolError::ProtocolVersionMismatch`
//!   plus clean shutdown afterwards.
//! * Message demultiplexer property: an arbitrary stream of valid
//!   envelopes (mixed responses, notifications, server requests) is
//!   correctly classified into the four buckets.
//!
//! ## Budget
//!
//! Every test in this file completes well under 5 seconds on a laptop.
//! Property tests are bounded to 32 cases per property by default; the
//! shapes are small so 32 is plenty without blowing past the budget.

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

use mcptest_core::protocol::{
    Client, ClientCapabilities, JsonRpcErrorBody, JsonRpcId, JsonRpcMessage, MessageKind,
    ProtocolError, ProtocolVersion, ServerRequestHandler, Transport,
};
use proptest::prelude::*;
use serde_json::{json, Value};
use tokio::sync::{mpsc, Mutex};
use tokio::time::timeout;

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// Build a mock transport pair plus a `MockPeer` the test can drive.
/// Mirrors the helper in `protocol_client.rs` so the two files can stay
/// independent (no shared `mod` import across integration tests).
fn mock_pair() -> (Transport, MockPeer) {
    let (out_tx, out_rx) = mpsc::unbounded_channel();
    let (in_tx, in_rx) = mpsc::unbounded_channel();
    (
        Transport::new(out_tx, in_rx),
        MockPeer {
            client_outbound: Arc::new(Mutex::new(out_rx)),
            client_inbound: in_tx,
        },
    )
}

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");
    }
}

fn build_runtime() -> tokio::runtime::Runtime {
    tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap()
}

// ---------------------------------------------------------------------------
// Handshake property: all four supported versions round trip
// ---------------------------------------------------------------------------

#[tokio::test]
async fn handshake_round_trips_each_supported_version() {
    for (client_version, wire) in [
        (ProtocolVersion::V2024_11_05, "2024-11-05"),
        (ProtocolVersion::V2025_03_26, "2025-03-26"),
        (ProtocolVersion::V2025_06_18, "2025-06-18"),
        (ProtocolVersion::V2026_03_26, "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;
            assert_eq!(req.method.as_deref(), Some("initialize"));
            let id = req.id.clone().unwrap();
            peer.send(JsonRpcMessage::response(
                id,
                json!({
                    "protocolVersion": wire,
                    "capabilities": {"tools": {}},
                    "serverInfo": {"name": "mock", "version": "0.1"}
                }),
            ));
            let note = peer.next_outbound().await;
            assert!(note.id.is_none(), "initialized must be a notification");
            assert_eq!(note.method.as_deref(), Some("initialized"));
            peer
        });
        let outcome = client.initialize(client_version).await.unwrap();
        assert_eq!(outcome.negotiated_version, client_version);
        assert_eq!(outcome.server.name.as_deref(), Some("mock"));
        let _peer = h.await.unwrap();
        client.close().await.unwrap();
    }
}

#[tokio::test]
async fn handshake_rejects_unknown_protocol_version_cleanly() {
    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": "2099-01-01"}),
        ));
        peer
    });
    let err = client
        .initialize(ProtocolVersion::current())
        .await
        .expect_err("must reject");
    assert!(matches!(err, ProtocolError::ProtocolVersionMismatch(ref v) if v == "2099-01-01"));
    let _peer = h.await.unwrap();
    // Transport tears down cleanly: close is idempotent and does not
    // hang even though the handshake aborted.
    client.close().await.unwrap();
}

// ---------------------------------------------------------------------------
// Capability negotiation: passthrough on both directions
// ---------------------------------------------------------------------------

#[tokio::test]
async fn client_capabilities_round_trip_through_initialize_request() {
    let caps = ClientCapabilities {
        name: "harness".into(),
        version: "9.9.9".into(),
        capabilities: json!({"sampling": {}, "roots": {"listChanged": true}}),
    };
    let (transport, peer) = mock_pair();
    let client = Client::builder(transport)
        .with_client_capabilities(caps.clone())
        .build();
    let h = tokio::spawn(async move {
        let req = peer.next_outbound().await;
        let params = req.params.clone().unwrap();
        assert_eq!(params["protocolVersion"], "2025-06-18");
        assert_eq!(params["clientInfo"]["name"], "harness");
        assert_eq!(params["clientInfo"]["version"], "9.9.9");
        assert_eq!(params["capabilities"]["sampling"], json!({}));
        assert_eq!(params["capabilities"]["roots"]["listChanged"], true);
        let id = req.id.clone().unwrap();
        peer.send(JsonRpcMessage::response(
            id,
            json!({"protocolVersion": "2025-06-18"}),
        ));
        let _note = peer.next_outbound().await;
        peer
    });
    let _outcome = client.initialize(ProtocolVersion::current()).await.unwrap();
    let _peer = h.await.unwrap();
    client.close().await.unwrap();
}

#[tokio::test]
async fn server_capabilities_arrive_intact_on_outcome() {
    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": "2025-06-18",
                "capabilities": {"tools": {"listChanged": true}, "resources": {}},
                "serverInfo": {"name": "echo", "version": "1.2.3"}
            }),
        ));
        let _note = peer.next_outbound().await;
        peer
    });
    let outcome = client.initialize(ProtocolVersion::current()).await.unwrap();
    assert_eq!(outcome.server.name.as_deref(), Some("echo"));
    assert_eq!(outcome.server.version.as_deref(), Some("1.2.3"));
    assert_eq!(outcome.server.capabilities["tools"]["listChanged"], true);
    assert!(outcome.server.capabilities.get("resources").is_some());
    let _peer = h.await.unwrap();
    client.close().await.unwrap();
}

// ---------------------------------------------------------------------------
// Error envelope handling: every canonical JSON-RPC code surfaces
// ---------------------------------------------------------------------------

async fn assert_server_error_surfaces(code: i64, msg: &str) {
    let (transport, peer) = mock_pair();
    let client = Client::builder(transport).build();
    let msg_owned = msg.to_owned();
    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,
                message: msg_owned,
                data: None,
            },
        ));
        peer
    });
    let err = client
        .request("tools/call", json!({}))
        .await
        .expect_err("must error");
    match err {
        ProtocolError::ServerError {
            code: c,
            message: m,
            ..
        } => {
            assert_eq!(c, code);
            assert_eq!(m, msg);
        }
        other => panic!("expected ServerError, got {other:?}"),
    }
    let _peer = h.await.unwrap();
    client.close().await.unwrap();
}

#[tokio::test]
async fn invalid_params_error_surfaces_as_typed_server_error() {
    assert_server_error_surfaces(-32602, "invalid params").await;
}

#[tokio::test]
async fn internal_error_surfaces_as_typed_server_error() {
    assert_server_error_surfaces(-32603, "internal error").await;
}

#[tokio::test]
async fn parse_error_surfaces_as_typed_server_error() {
    assert_server_error_surfaces(-32700, "parse error").await;
}

#[tokio::test]
async fn application_error_surfaces_as_typed_server_error() {
    assert_server_error_surfaces(-32000, "application error").await;
}

#[tokio::test]
async fn error_data_payload_passes_through() {
    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: -32602,
                message: "invalid params".into(),
                data: Some(json!({"field": "name", "reason": "missing"})),
            },
        ));
        peer
    });
    let err = client
        .request("tools/call", json!({}))
        .await
        .expect_err("must error");
    match err {
        ProtocolError::ServerError { data, .. } => {
            let data = data.expect("data payload preserved");
            assert_eq!(data["field"], "name");
            assert_eq!(data["reason"], "missing");
        }
        other => panic!("expected ServerError, got {other:?}"),
    }
    let _peer = h.await.unwrap();
    client.close().await.unwrap();
}

// ---------------------------------------------------------------------------
// Notification routing: handler receives every notification
// ---------------------------------------------------------------------------

#[tokio::test]
async fn notifications_route_to_handler_and_skip_inflight_map() {
    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);
        tokio::spawn(async move {
            s.lock().await.push((method, params));
        });
    });
    let client = Client::builder(transport)
        .with_notification_handler(handler)
        .build();
    for n in 0..5 {
        peer.send(JsonRpcMessage::notification(
            "notifications/progress",
            json!({"value": n}),
        ));
    }
    tokio::time::sleep(Duration::from_millis(80)).await;
    let collected = sink.lock().await.clone();
    assert_eq!(collected.len(), 5);
    for (i, (method, params)) in collected.iter().enumerate() {
        assert_eq!(method, "notifications/progress");
        assert_eq!(params["value"], i);
    }
    drop(peer);
    client.close().await.unwrap();
}

// ---------------------------------------------------------------------------
// Server-initiated request handling routes to the registered handler
// ---------------------------------------------------------------------------

#[tokio::test]
async fn server_initiated_request_routes_to_handler_and_echoes_back() {
    let (transport, peer) = mock_pair();
    let handler: 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(99),
        "sampling/createMessage",
        json!({"prompt": "hi"}),
    ));
    let reply = timeout(Duration::from_millis(500), peer.next_outbound())
        .await
        .expect("client replied");
    assert_eq!(reply.id, Some(JsonRpcId::Number(99)));
    let result = reply.result.expect("success response");
    assert_eq!(result["echo"]["prompt"], "hi");
    drop(peer);
    client.close().await.unwrap();
}

// ---------------------------------------------------------------------------
// Concurrent in-flight correlation under proptest permutations
// ---------------------------------------------------------------------------

proptest! {
    #![proptest_config(ProptestConfig {
        cases: 8,
        ..ProptestConfig::default()
    })]

    /// Property: dispatch N requests concurrently, respond in an
    /// arbitrary permutation, and assert every caller receives the
    /// response that corresponds to its own request ID. The runner does
    /// not own the Tokio reactor here; proptest is sync, so each case
    /// builds a current-thread runtime and drives the case to
    /// completion. N is bounded at 10 to keep the case quick.
    #[test]
    fn concurrent_inflight_requests_correlate_under_any_response_order(
        perm in proptest::collection::vec(0u8..10, 10).prop_filter("unique permutation", |v| {
            let mut seen = std::collections::HashSet::new();
            v.iter().all(|x| seen.insert(*x))
        })
    ) {
        let rt = build_runtime();
        rt.block_on(async move {
            let (transport, peer) = mock_pair();
            let client = Client::builder(transport)
                .with_request_timeout(Duration::from_secs(5))
                .build();
            // Drain N requests, then reply in the permuted order. Each
            // response carries the index as its `who` so the assertion
            // below can cross-check on payload.
            let peer = Arc::new(peer);
            let peer_for_task = Arc::clone(&peer);
            let perm_for_task = perm.clone();
            let peer_handle = tokio::spawn(async move {
                let mut reqs = Vec::with_capacity(10);
                for _ in 0..10 {
                    reqs.push(peer_for_task.next_outbound().await);
                }
                for idx in perm_for_task.iter() {
                    let req = &reqs[*idx as usize];
                    peer_for_task.send(JsonRpcMessage::response(
                        req.id.clone().unwrap(),
                        json!({"who": *idx}),
                    ));
                }
            });
            let mut handles = Vec::with_capacity(10);
            for idx in 0..10u8 {
                let c = client.clone();
                handles.push(tokio::spawn(async move {
                    let r = c.request("tools/list", json!({"who": idx})).await.unwrap();
                    (idx, r)
                }));
            }
            for h in handles {
                let (idx, r) = h.await.unwrap();
                let who = r.get("who").and_then(Value::as_u64).unwrap() as u8;
                assert_eq!(who, idx, "request {idx} got response for {who}");
            }
            let _ = peer_handle.await;
            client.close().await.unwrap();
        });
    }
}

// ---------------------------------------------------------------------------
// Message demultiplexer property: arbitrary envelopes classify correctly
// ---------------------------------------------------------------------------

#[derive(Debug, Clone)]
enum Sample {
    Request(i64, String),
    Response(i64),
    ErrorResponse(i64, i64),
    Notification(String),
}

fn sample() -> impl Strategy<Value = Sample> {
    prop_oneof![
        (any::<i64>(), "[a-z]{1,8}").prop_map(|(id, m)| Sample::Request(id, m)),
        any::<i64>().prop_map(Sample::Response),
        (any::<i64>(), -32700_i64..-32000).prop_map(|(id, code)| Sample::ErrorResponse(id, code)),
        "[a-z]{1,8}".prop_map(Sample::Notification),
    ]
}

proptest! {
    #![proptest_config(ProptestConfig {
        cases: 64,
        ..ProptestConfig::default()
    })]

    /// Property: every well-formed envelope shape classifies into the
    /// matching `MessageKind` variant. The classifier is the single
    /// source of truth for which bucket the inbound pump uses, so a
    /// regression here would silently break correlation, notification
    /// routing, or server-request dispatch.
    #[test]
    fn envelope_kind_classifier_matches_construction(s in sample()) {
        let msg = match s.clone() {
            Sample::Request(id, m) => JsonRpcMessage::request(
                JsonRpcId::Number(id), m, json!({}),
            ),
            Sample::Response(id) => JsonRpcMessage::response(
                JsonRpcId::Number(id), json!({}),
            ),
            Sample::ErrorResponse(id, code) => JsonRpcMessage::error_response(
                JsonRpcId::Number(id),
                JsonRpcErrorBody { code, message: "x".into(), data: None },
            ),
            Sample::Notification(m) => JsonRpcMessage::notification(m, json!({})),
        };
        let kind = msg.kind().expect("well-formed envelope classifies");
        match (s, kind) {
            (Sample::Request(_, _), MessageKind::Request { .. }) => {},
            (Sample::Response(_), MessageKind::Response { .. }) => {},
            (Sample::ErrorResponse(_, _), MessageKind::ErrorResponse { .. }) => {},
            (Sample::Notification(_), MessageKind::Notification { .. }) => {},
            (s, k) => prop_assert!(false, "sample {s:?} classified as {k:?}"),
        }
    }
}

// ---------------------------------------------------------------------------
// Malformed envelopes return ProtocolError::MalformedMessage, not a panic
// ---------------------------------------------------------------------------

#[test]
fn envelope_with_both_result_and_error_is_rejected() {
    let m = JsonRpcMessage {
        jsonrpc: "2.0".into(),
        id: Some(JsonRpcId::Number(1)),
        method: None,
        params: None,
        result: Some(json!({})),
        error: Some(JsonRpcErrorBody {
            code: -32601,
            message: "x".into(),
            data: None,
        }),
    };
    assert!(matches!(m.kind(), Err(ProtocolError::MalformedMessage(_))));
}

#[test]
fn envelope_with_wrong_jsonrpc_version_is_rejected() {
    let m = JsonRpcMessage {
        jsonrpc: "1.0".into(),
        id: Some(JsonRpcId::Number(1)),
        method: Some("ping".into()),
        params: None,
        result: None,
        error: None,
    };
    assert!(matches!(m.kind(), Err(ProtocolError::MalformedMessage(_))));
}
