//! Cassette recording transport.
//!
//! Wraps a live transport so a run can capture what really happened on the
//! wire into a [`Cassette`] that [`ReplayTransport`](super::ReplayTransport)
//! can replay later. The wrapper sits between the protocol [`Client`] and the
//! real transport's channels: it forwards every message untouched and tees
//! request/response pairs into the recording.
//!
//! ## Pairing
//!
//! A request carries an id; its response echoes that id. The forwarding task
//! remembers each outbound request by id and, when the matching response
//! arrives, emits an [`Exchange`]. Notifications (no id) are forwarded but not
//! recorded, since they have no response and the replayer swallows them.
//! Exchanges are stored in response-arrival order, which for the sequential
//! request/await client is request order.

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use tokio::sync::mpsc;
use tokio::task::JoinHandle;

use crate::cassette::{Cassette, Exchange};
use crate::protocol::{JsonRpcId, JsonRpcMessage, Transport};

/// Shared state the forwarding task writes and [`RecordingHandle::finish`]
/// reads once the task has stopped.
#[derive(Default, Debug)]
struct RecordState {
    /// Outbound requests awaiting a response, keyed by id.
    pending: HashMap<JsonRpcId, JsonRpcMessage>,
    /// Completed request/response pairs in arrival order.
    exchanges: Vec<Exchange>,
}

/// Handle to a recording in progress. Call [`RecordingHandle::finish`] to stop
/// forwarding and take the captured cassette.
#[derive(Debug)]
pub struct RecordingHandle {
    state: Arc<Mutex<RecordState>>,
    task: JoinHandle<()>,
}

impl RecordingHandle {
    /// Copy the exchanges captured so far into a cassette without
    /// stopping the tap.
    ///
    /// Exists for the proxy recorder (`mcptest record`), which rewrites
    /// the cassette file after every completed exchange so a crash or an
    /// abrupt client disconnect never loses finished work. The forwarding
    /// task keeps running; only the clone happens under the lock.
    pub fn snapshot(&self) -> Cassette {
        let state = self.state.lock().expect("record state not poisoned");
        let mut cassette = Cassette::new();
        for exchange in &state.exchanges {
            cassette.push(exchange.clone());
        }
        cassette
    }

    /// Stop forwarding and return the recorded cassette.
    ///
    /// The cassette is returned un-redacted; the caller is expected to run it
    /// through [`Cassette::redacted`] before writing it to disk so secrets in
    /// recorded params never land in a file.
    pub async fn finish(self) -> Cassette {
        self.task.abort();
        let _ = self.task.await;
        let state = self.state.lock().expect("record state not poisoned");
        let mut cassette = Cassette::new();
        for exchange in &state.exchanges {
            cassette.push(exchange.clone());
        }
        cassette
    }
}

/// Recording wrapper around a live transport.
#[derive(Debug)]
pub struct RecordingTransport;

impl RecordingTransport {
    /// Interpose a recorder between the client and `inner`. Returns the
    /// transport the client should drive plus the handle that yields the
    /// cassette once the run is done.
    pub fn wrap(inner: Transport) -> (Transport, RecordingHandle) {
        let Transport {
            outbound_tx: real_out,
            inbound_rx: mut real_in,
        } = inner;
        let (client_out_tx, mut client_out_rx) = mpsc::unbounded_channel::<JsonRpcMessage>();
        let (client_in_tx, client_in_rx) = mpsc::unbounded_channel::<JsonRpcMessage>();
        let state = Arc::new(Mutex::new(RecordState::default()));
        let task_state = Arc::clone(&state);

        let task = tokio::spawn(async move {
            loop {
                tokio::select! {
                    out = client_out_rx.recv() => {
                        let Some(out) = out else { break };
                        if let Some(id) = out.id.clone() {
                            task_state
                                .lock()
                                .expect("record state")
                                .pending
                                .insert(id, out.clone());
                        }
                        if real_out.send(out).is_err() {
                            break;
                        }
                    }
                    inbound = real_in.recv() => {
                        let Some(inbound) = inbound else { break };
                        if let Some(id) = inbound.id.clone() {
                            let mut guard = task_state.lock().expect("record state");
                            if let Some(request) = guard.pending.remove(&id) {
                                guard.exchanges.push(Exchange {
                                    request,
                                    response: inbound.clone(),
                                });
                            }
                        }
                        if client_in_tx.send(inbound).is_err() {
                            break;
                        }
                    }
                }
            }
        });

        (
            Transport::new(client_out_tx, client_in_rx),
            RecordingHandle { state, task },
        )
    }
}

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

    #[tokio::test]
    async fn records_a_request_response_pair_and_forwards_both_ways() {
        // Build the "real server" side of the channels by hand so the test can
        // see what the recorder forwarded and inject a response.
        let (real_out_tx, mut server_rx) = mpsc::unbounded_channel::<JsonRpcMessage>();
        let (server_tx, real_in_rx) = mpsc::unbounded_channel::<JsonRpcMessage>();
        let inner = Transport::new(real_out_tx, real_in_rx);

        let (client, handle) = RecordingTransport::wrap(inner);
        let Transport {
            outbound_tx: client_out,
            inbound_rx: mut client_in,
        } = client;

        // Client sends a request; the recorder forwards it to the server.
        client_out
            .send(JsonRpcMessage::request(
                JsonRpcId::Number(5),
                "tools/list",
                json!({}),
            ))
            .unwrap();
        let forwarded = server_rx.recv().await.expect("forwarded to server");
        assert_eq!(forwarded.method.as_deref(), Some("tools/list"));

        // Server responds; the recorder forwards it back to the client.
        server_tx
            .send(JsonRpcMessage::response(
                JsonRpcId::Number(5),
                json!({"tools": [{"name": "search"}]}),
            ))
            .unwrap();
        let got = client_in.recv().await.expect("response to client");
        assert_eq!(got.id, Some(JsonRpcId::Number(5)));

        let cassette = handle.finish().await;
        assert_eq!(cassette.exchanges.len(), 1);
        assert_eq!(
            cassette.exchanges[0].request.method.as_deref(),
            Some("tools/list")
        );
        assert_eq!(
            cassette.exchanges[0].response.result.as_ref().unwrap()["tools"][0]["name"],
            "search"
        );
    }

    #[tokio::test]
    async fn snapshot_returns_exchanges_so_far_without_stopping_the_tap() {
        let (real_out_tx, mut server_rx) = mpsc::unbounded_channel::<JsonRpcMessage>();
        let (server_tx, real_in_rx) = mpsc::unbounded_channel::<JsonRpcMessage>();
        let inner = Transport::new(real_out_tx, real_in_rx);
        let (client, handle) = RecordingTransport::wrap(inner);
        let Transport {
            outbound_tx: client_out,
            inbound_rx: mut client_in,
        } = client;

        // First exchange completes, then a snapshot is taken mid-recording.
        client_out
            .send(JsonRpcMessage::request(
                JsonRpcId::Number(1),
                "tools/list",
                json!({}),
            ))
            .unwrap();
        let _ = server_rx.recv().await.unwrap();
        server_tx
            .send(JsonRpcMessage::response(JsonRpcId::Number(1), json!({})))
            .unwrap();
        let _ = client_in.recv().await.unwrap();
        let snapshot = handle.snapshot();
        assert_eq!(snapshot.exchanges.len(), 1);

        // The tap is still live: a second exchange flows and lands in the
        // final cassette alongside the first.
        client_out
            .send(JsonRpcMessage::request(
                JsonRpcId::Number(2),
                "tools/call",
                json!({"name": "x"}),
            ))
            .unwrap();
        let _ = server_rx.recv().await.unwrap();
        server_tx
            .send(JsonRpcMessage::response(JsonRpcId::Number(2), json!({})))
            .unwrap();
        let _ = client_in.recv().await.unwrap();
        let cassette = handle.finish().await;
        assert_eq!(cassette.exchanges.len(), 2);
        assert_eq!(
            snapshot.exchanges.len(),
            1,
            "snapshot is a point-in-time copy"
        );
    }

    #[tokio::test]
    async fn notifications_are_forwarded_but_not_recorded() {
        let (real_out_tx, mut server_rx) = mpsc::unbounded_channel::<JsonRpcMessage>();
        let (_server_tx, real_in_rx) = mpsc::unbounded_channel::<JsonRpcMessage>();
        let inner = Transport::new(real_out_tx, real_in_rx);
        let (client, handle) = RecordingTransport::wrap(inner);
        let Transport {
            outbound_tx: client_out,
            ..
        } = client;

        client_out
            .send(JsonRpcMessage::notification(
                "notifications/initialized",
                json!({}),
            ))
            .unwrap();
        let forwarded = server_rx.recv().await.expect("notification forwarded");
        assert_eq!(forwarded.id, None);

        let cassette = handle.finish().await;
        assert!(
            cassette.exchanges.is_empty(),
            "a notification has no response to record"
        );
    }

    #[tokio::test]
    async fn recorded_cassette_redacts_secrets_on_save() {
        let (real_out_tx, mut server_rx) = mpsc::unbounded_channel::<JsonRpcMessage>();
        let (server_tx, real_in_rx) = mpsc::unbounded_channel::<JsonRpcMessage>();
        let inner = Transport::new(real_out_tx, real_in_rx);
        let (client, handle) = RecordingTransport::wrap(inner);
        let Transport {
            outbound_tx: client_out,
            inbound_rx: mut client_in,
        } = client;

        client_out
            .send(JsonRpcMessage::request(
                JsonRpcId::Number(1),
                "tools/call",
                json!({"name": "x", "arguments": {"api_key": "sk-abc1234567890DEF"}}),
            ))
            .unwrap();
        let _ = server_rx.recv().await.unwrap();
        server_tx
            .send(JsonRpcMessage::response(
                JsonRpcId::Number(1),
                json!({"ok": true}),
            ))
            .unwrap();
        let _ = client_in.recv().await.unwrap();

        let cassette = handle.finish().await;
        let redacted = cassette.redacted(&crate::redaction::Redactor::new());
        let key = &redacted.exchanges[0].request.params.as_ref().unwrap()["arguments"]["api_key"];
        assert!(
            key.as_str().unwrap().starts_with("[REDACTED:"),
            "secret should be scrubbed before save, got {key}"
        );
    }
}
