//! Agent-cassette extraction for distillation.
//!
//! Why this exists: an agent cassette (one recorded model loop, see
//! `schemas/agent-cassette-v1.json`) carries the tool legs the model
//! actually drove: each call's name, args, and the verbatim MCP result.
//! Those legs distill into tool tests the same way server exchanges do.
//! The cassette names the suite's server *key* per call but not how to
//! reach that server (no command or URL is recorded), so the generated
//! `servers:` block is a clearly-commented template the user retargets.
//!
//! The types here are a deliberately minimal mirror of the cassette
//! schema rather than a dependency on the `mcptest-agent` crate: core
//! sits below the agent crate in the workspace graph, and distillation
//! only needs the tool legs.

use serde::Deserialize;
use serde_json::{json, Value};

use super::{Extraction, RecordedCall};
use crate::redaction::Redactor;

/// The slice of an agent cassette distillation reads. Unknown fields
/// (prompts, conversation summary, timestamps) are ignored on parse since
/// only the tool legs become tests.
#[derive(Debug, Clone, Deserialize)]
pub struct AgentRecording {
    /// Agent `name:` from the YAML suite that recorded the cassette.
    pub agent_name: String,
    /// Model identifier the recording used. Diagnostic only here.
    pub model: String,
    /// The recorded conversation's tool legs.
    pub trace: AgentTrace,
}

/// Tool legs of one recorded agent conversation.
#[derive(Debug, Clone, Deserialize)]
pub struct AgentTrace {
    /// Every tool call the model issued, in order.
    #[serde(default)]
    pub tool_calls: Vec<AgentToolCall>,
    /// The MCP results, linked back to calls via `call_index`.
    #[serde(default)]
    pub tool_results: Vec<AgentToolResult>,
}

/// One tool call the model issued during the recording.
#[derive(Debug, Clone, Deserialize)]
pub struct AgentToolCall {
    /// Tool name without the `<server>__` namespace prefix.
    pub name: String,
    /// Suite server key the call routed to. Empty on cassettes recorded
    /// before the field existed, in which case the caller's default
    /// server name is used.
    #[serde(default)]
    pub server: String,
    /// Arguments the model passed, free-form JSON.
    #[serde(default)]
    pub args: Value,
}

/// One MCP result captured during the recording.
#[derive(Debug, Clone, Deserialize)]
pub struct AgentToolResult {
    /// Index of the matching entry in `tool_calls`.
    pub call_index: usize,
    /// Verbatim JSON-RPC `result` payload from the MCP server.
    pub result: Value,
}

/// Lift every completed tool leg (call plus result) out of an agent
/// recording. Args and results pass through the [`Redactor`] before they
/// can reach YAML; calls without a recorded result are noted and skipped.
pub(crate) fn extract(recording: &AgentRecording, default_server: &str) -> Extraction {
    let redactor = Redactor::new();
    let mut calls = Vec::new();
    let mut skipped = Vec::new();
    for (index, call) in recording.trace.tool_calls.iter().enumerate() {
        let Some(result) = recording
            .trace
            .tool_results
            .iter()
            .find(|r| r.call_index == index)
        else {
            skipped.push(format!("call #{index} {} (no recorded result)", call.name));
            continue;
        };
        let args = match &call.args {
            Value::Object(_) => redactor.redact_json(&call.args),
            // tools/call arguments are an object on the wire; anything else
            // cannot render as a schema-valid `args:` block, so start the
            // test from empty args and let the skip note say why.
            other => {
                if !other.is_null() {
                    skipped.push(format!(
                        "call #{index} {}: non-object args were replaced with {{}}",
                        call.name
                    ));
                }
                json!({})
            }
        };
        let server = if call.server.is_empty() {
            default_server.to_string()
        } else {
            call.server.clone()
        };
        calls.push(RecordedCall {
            server,
            tool: call.name.clone(),
            args,
            result: redactor.redact_json(&result.result),
            source_index: index,
        });
    }
    let notes = vec![format!(
        "agent cassettes (agent `{}`, model `{}`) record which server key served each \
         call but not how to reach it; the generated servers: block is a template you \
         must retarget at the real server before running",
        recording.agent_name, recording.model
    )];
    Extraction {
        servers_block: servers_block(&calls, default_server),
        notes,
        calls,
        skipped,
    }
}

/// Render the template `servers:` block for the server keys the recording
/// names. Each entry is schema-valid placeholder `command:` so the suite
/// validates as written, with a TODO comment telling the user to retarget.
fn servers_block(calls: &[RecordedCall], default_server: &str) -> String {
    let mut names: Vec<&str> = calls.iter().map(|c| c.server.as_str()).collect();
    if names.is_empty() {
        names.push(default_server);
    }
    names.sort_unstable();
    names.dedup();
    let mut out = String::from("servers:\n");
    out.push_str(
        "  # TODO(retarget): the agent cassette names server keys but not how to\n  \
         # reach them. Replace each placeholder command with the real server entry\n  \
         # from the suite the agent ran against.\n",
    );
    for name in names {
        out.push_str(&format!("  {name}:\n"));
        out.push_str("    command: [\"./path/to/your-mcp-server\"]\n");
    }
    out
}

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

    fn recording(calls: Value, results: Value) -> AgentRecording {
        serde_json::from_value(json!({
            "agent_name": "weather query",
            "model": "claude-sonnet-4-5",
            "trace": { "tool_calls": calls, "tool_results": results }
        }))
        .expect("recording parses")
    }

    #[test]
    fn pairs_calls_with_results_by_call_index_and_keeps_server_keys() {
        let rec = recording(
            json!([
                { "name": "get_weather", "server": "weather", "args": { "city": "Sacramento" } },
                { "name": "get_time", "server": "clock", "args": {} }
            ]),
            json!([
                { "call_index": 1, "is_error": false, "result": { "content": [] } },
                { "call_index": 0, "is_error": false, "result": { "content": [] } }
            ]),
        );
        let ex = extract(&rec, "target");
        assert_eq!(ex.calls.len(), 2);
        assert_eq!(ex.calls[0].tool, "get_weather");
        assert_eq!(ex.calls[0].server, "weather");
        assert_eq!(ex.calls[0].args, json!({ "city": "Sacramento" }));
        assert_eq!(ex.calls[1].server, "clock");
        // Both named servers get a template entry.
        assert!(ex.servers_block.contains("  weather:\n"));
        assert!(ex.servers_block.contains("  clock:\n"));
        assert!(ex.servers_block.contains("TODO(retarget)"));
    }

    #[test]
    fn call_without_a_recorded_result_is_skipped_with_a_reason() {
        let rec = recording(json!([{ "name": "lonely", "args": {} }]), json!([]));
        let ex = extract(&rec, "target");
        assert!(ex.calls.is_empty());
        assert_eq!(ex.skipped, vec!["call #0 lonely (no recorded result)"]);
    }

    #[test]
    fn empty_server_key_falls_back_to_the_default_server_name() {
        let rec = recording(
            json!([{ "name": "echo", "args": {} }]),
            json!([{ "call_index": 0, "is_error": false, "result": { "content": [] } }]),
        );
        let ex = extract(&rec, "target");
        assert_eq!(ex.calls[0].server, "target");
        assert!(ex.servers_block.contains("  target:\n"));
    }

    #[test]
    fn non_object_args_become_empty_args_with_a_note() {
        let rec = recording(
            json!([{ "name": "echo", "args": [1, 2] }]),
            json!([{ "call_index": 0, "is_error": false, "result": { "content": [] } }]),
        );
        let ex = extract(&rec, "target");
        assert_eq!(ex.calls[0].args, json!({}));
        assert!(ex.skipped.iter().any(|s| s.contains("non-object args")));
    }
}
