//! Distill a recorded cassette into a draft YAML suite (WOR-1263).
//!
//! Why this exists: recording already works (server cassettes capture
//! JSON-RPC exchanges, agent cassettes capture a model loop's tool legs),
//! but a recording is not a suite. Nobody converts real traffic into
//! assertion-level regression tests by hand. This module is that
//! converter: each recorded `tools/call` becomes one test whose args are
//! lifted verbatim and whose assertions come from [`crate::propose`], so
//! distillation and live proposal share one derivation policy.
//!
//! Single-observation semantics apply: a call seen once gets the
//! `isError` invariant and the structural schema only. When the same
//! tool+args pair was recorded more than once, the repeats act as the
//! stability probe and equal leaves earn exact-value assertions.
//!
//! Recorded traffic is real traffic, so two safety layers run before any
//! value reaches YAML: the shared [`crate::redaction`] pass scrubs
//! secret-shaped strings, and the review heuristics flag what survives
//! but still looks personal or live (emails, bearer-like tokens,
//! non-example URLs). Flags never rewrite; they add a `# review:`
//! comment above the affected test.

mod agent;
mod review;
mod server;

use std::collections::BTreeSet;
use std::path::Path;
use std::time::Duration;

use serde_json::Value;

use crate::cassette::Cassette;
use crate::propose::{propose_test, Observation, ProposalInput};

pub use agent::{AgentRecording, AgentToolCall, AgentToolResult, AgentTrace};
pub use review::{ReviewCategory, ReviewFlag};

/// Comment line placed directly above a test whose recorded args or
/// asserted exact values tripped a review heuristic.
pub const REVIEW_COMMENT: &str = "# review: possible personal or live data";

/// Errors raised while parsing or distilling a cassette.
///
/// Typed so the CLI can attach the cassette path and the operator's fix
/// is obvious per variant (fix the JSON, record some tool calls).
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum DistillError {
    /// The file is not valid JSON, or it matched a known cassette shape
    /// but failed to deserialize as one.
    #[error("could not parse cassette: {0}")]
    Parse(#[from] serde_json::Error),
    /// The JSON parses but matches neither cassette shape.
    #[error(
        "unrecognized cassette shape: expected a server cassette (with `exchanges`) or an \
         agent cassette (with `trace`)"
    )]
    UnknownShape,
    /// The cassette holds no `tools/call` exchange (or no completed agent
    /// tool leg), so there is nothing to turn into a test.
    #[error("the cassette contains no completed tools/call exchanges to distill")]
    NoToolCalls,
}

/// A cassette distillation understands, discriminated by shape.
///
/// Server cassettes are ordered JSON-RPC exchanges; agent cassettes are
/// one recorded model loop whose tool legs carry the same `tools/call`
/// results. Both distill through the same proposal pipeline.
#[derive(Debug, Clone)]
pub enum ParsedCassette {
    /// A server cassette (see [`crate::cassette`]).
    Server(Cassette),
    /// An agent cassette (see `schemas/agent-cassette-v1.json`).
    Agent(AgentRecording),
}

/// Parse raw cassette JSON, detecting which kind it is by shape: a
/// top-level `exchanges` array means a server cassette, a top-level
/// `trace` object means an agent cassette.
pub fn parse_cassette(raw: &str) -> Result<ParsedCassette, DistillError> {
    let value: Value = serde_json::from_str(raw)?;
    if value.get("exchanges").is_some() {
        return Ok(ParsedCassette::Server(serde_json::from_value(value)?));
    }
    if value.get("trace").is_some() {
        return Ok(ParsedCassette::Agent(serde_json::from_value(value)?));
    }
    Err(DistillError::UnknownShape)
}

/// How the distilled suite should reference its source.
#[derive(Debug, Clone)]
pub struct DistillOptions {
    /// Cassette path as the suite should reference it (the `cassette:`
    /// value for server cassettes; also used in header comments and to
    /// derive the test-name stem).
    pub cassette_path: String,
    /// Server key for the generated `servers:` block and for tests whose
    /// recording does not name a server itself.
    pub server_name: String,
}

/// Metadata for one generated test, so callers can summarize the suite
/// without re-parsing the YAML.
#[derive(Debug, Clone)]
pub struct DistilledTest {
    /// Final (deduplicated) test name as it appears in the YAML.
    pub name: String,
    /// Server key the test runs against.
    pub server: String,
    /// Tool the test calls.
    pub tool: String,
    /// Where the call came from: the exchange index (server cassettes)
    /// or the tool-call index (agent cassettes).
    pub source_index: usize,
}

/// A distilled draft suite: the complete YAML document plus the metadata
/// and review flags the CLI surfaces on stderr.
#[derive(Debug, Clone)]
pub struct DistilledSuite {
    /// Complete suite document: schema header, `servers:` block, and one
    /// test per recorded tool call under `tools:`.
    pub yaml: String,
    /// One entry per generated test, in YAML order.
    pub tests: Vec<DistilledTest>,
    /// Suspicious recorded values (emails, bearer-like tokens, live
    /// URLs) found in args or exact assertions. Flag-only; the YAML
    /// values are untouched.
    pub review_flags: Vec<ReviewFlag>,
    /// Human-readable caveats (replay-order warnings, agent retarget
    /// note) for the CLI to print to stderr.
    pub notes: Vec<String>,
}

/// One recorded tools/call leg, normalized across both cassette kinds.
#[derive(Debug, Clone)]
pub(crate) struct RecordedCall {
    /// Server key the generated test references.
    pub server: String,
    /// Tool name.
    pub tool: String,
    /// Arguments lifted verbatim (post-redaction) from the recording.
    pub args: Value,
    /// The recorded `tools/call` result (post-redaction).
    pub result: Value,
    /// Exchange index or tool-call index in the source cassette.
    pub source_index: usize,
}

/// What an extraction pass produced: the calls to distill, the skip
/// summary, stderr-level notes, and the rendered `servers:` block.
#[derive(Debug, Clone)]
pub(crate) struct Extraction {
    pub calls: Vec<RecordedCall>,
    pub skipped: Vec<String>,
    pub notes: Vec<String>,
    pub servers_block: String,
}

/// Distill a parsed cassette into a draft suite.
///
/// One test per recorded `tools/call`, args verbatim, assertions via
/// [`crate::propose`]. A tool+args pair recorded more than once uses its
/// repeats as the stability probe (exact-value assertions become
/// eligible); a pair seen once keeps single-observation semantics.
///
/// ```
/// use mcptest_core::distill::{distill, parse_cassette, DistillOptions};
///
/// let raw = r#"{ "version": "1", "exchanges": [
///   { "request": { "jsonrpc": "2.0", "id": 1, "method": "tools/call",
///       "params": { "name": "echo", "arguments": { "msg": "hi" } } },
///     "response": { "jsonrpc": "2.0", "id": 1,
///       "result": { "content": [{ "type": "text", "text": "hi" }] } } }
/// ]}"#;
/// let cassette = parse_cassette(raw).expect("parses");
/// let suite = distill(
///     &cassette,
///     &DistillOptions {
///         cassette_path: "./tape.json".to_string(),
///         server_name: "recorded".to_string(),
///     },
/// )
/// .expect("distills");
/// assert!(suite.yaml.contains("echo: distilled from tape #1"));
/// ```
pub fn distill(
    cassette: &ParsedCassette,
    opts: &DistillOptions,
) -> Result<DistilledSuite, DistillError> {
    let extraction = match cassette {
        ParsedCassette::Server(c) => server::extract(c, &opts.server_name, &opts.cassette_path),
        ParsedCassette::Agent(a) => agent::extract(a, &opts.server_name),
    };
    if extraction.calls.is_empty() {
        return Err(DistillError::NoToolCalls);
    }
    let stem = cassette_stem(&opts.cassette_path);
    let mut used_names = BTreeSet::new();
    let mut blocks = Vec::new();
    let mut tests = Vec::new();
    let mut review_flags = Vec::new();
    for (index, call) in extraction.calls.iter().enumerate() {
        let proposal = propose_for(index, &extraction.calls);
        let name = unique_name(
            format!("{}: distilled from {} #{}", call.tool, stem, index + 1),
            &mut used_names,
        );
        let mut block = rename_block(&proposal.yaml, &name);
        let flags = scan_test(&name, call, &block);
        if !flags.is_empty() {
            block = format!("  {REVIEW_COMMENT}\n{block}");
        }
        review_flags.extend(flags);
        blocks.push(block);
        tests.push(DistilledTest {
            name,
            server: call.server.clone(),
            tool: call.tool.clone(),
            source_index: call.source_index,
        });
    }
    Ok(DistilledSuite {
        yaml: render_suite(opts, &extraction, &blocks),
        tests,
        review_flags,
        notes: extraction.notes,
    })
}

/// Derive the proposal for one call. The recording carries no timing, so
/// durations are zero and the proposed `max_duration_ms` field sits at the
/// floor, a deliberately generous budget for a replayed exchange.
fn propose_for(index: usize, calls: &[RecordedCall]) -> crate::propose::ProposedTest {
    let call = &calls[index];
    let obs = |c: &RecordedCall| Observation {
        result: c.result.clone(),
        duration: Duration::ZERO,
    };
    let partner = stability_partner(index, calls).map(|i| obs(&calls[i]));
    let input = ProposalInput {
        server: call.server.clone(),
        tool: call.tool.clone(),
        args: call.args.clone(),
    };
    propose_test(&input, &obs(call), partner.as_ref())
}

/// Find another recording of the same tool+args pair to act as the
/// stability probe. Returns the next repeat in recorded order (wrapping),
/// or `None` when the pair was seen once.
fn stability_partner(index: usize, calls: &[RecordedCall]) -> Option<usize> {
    let key = group_key(&calls[index]);
    let positions: Vec<usize> = calls
        .iter()
        .enumerate()
        .filter(|(_, c)| group_key(c) == key)
        .map(|(i, _)| i)
        .collect();
    if positions.len() < 2 {
        return None;
    }
    let pos = positions.iter().position(|&p| p == index)?;
    Some(positions[(pos + 1) % positions.len()])
}

/// Canonical identity of a call: server, tool, and the compact JSON form
/// of its args (serde_json maps are sorted, so equal objects render
/// identically). Unit separator avoids accidental collisions.
fn group_key(call: &RecordedCall) -> String {
    format!("{}\u{1f}{}\u{1f}{}", call.server, call.tool, call.args)
}

/// Filename stem of the cassette, used in generated test names.
fn cassette_stem(path: &str) -> String {
    Path::new(path)
        .file_stem()
        .and_then(|s| s.to_str())
        .unwrap_or("cassette")
        .to_string()
}

/// Make `base` unique by appending ` (2)`, ` (3)`, ... on collision.
/// Ordinal-numbered names never collide in practice; this is the
/// guarantee that they cannot.
fn unique_name(base: String, used: &mut BTreeSet<String>) -> String {
    if used.insert(base.clone()) {
        return base;
    }
    let mut k = 2usize;
    loop {
        let candidate = format!("{base} ({k})");
        if used.insert(candidate.clone()) {
            return candidate;
        }
        k += 1;
    }
}

/// Replace the proposal's generated `- name:` line with the distilled
/// name. The propose renderer always emits the name first; if that ever
/// changes the block is kept intact rather than corrupted.
fn rename_block(block: &str, name: &str) -> String {
    match block.split_once('\n') {
        Some((first, rest)) if first.trim_start().starts_with("- name:") => {
            format!("  - name: {}\n{rest}", yaml_quote(name))
        }
        _ => block.to_string(),
    }
}

/// Scan one generated test for review-worthy values: every string leaf of
/// the recorded args, plus every exact-value assertion string in the
/// rendered block. Flag-only; nothing is rewritten.
fn scan_test(test_name: &str, call: &RecordedCall, block: &str) -> Vec<ReviewFlag> {
    let mut flags = Vec::new();
    scan_value(test_name, "args", &call.args, &mut flags);
    for (target, value) in exact_assertion_strings(block) {
        if let Some(category) = review::categorize(&value) {
            flags.push(ReviewFlag {
                test: test_name.to_string(),
                path: target,
                category,
            });
        }
    }
    flags
}

/// Walk a JSON value and flag suspicious string leaves, building
/// `args.user.email` style paths.
fn scan_value(test_name: &str, prefix: &str, value: &Value, out: &mut Vec<ReviewFlag>) {
    match value {
        Value::String(s) => {
            if let Some(category) = review::categorize(s) {
                out.push(ReviewFlag {
                    test: test_name.to_string(),
                    path: prefix.to_string(),
                    category,
                });
            }
        }
        Value::Array(items) => {
            for (i, item) in items.iter().enumerate() {
                scan_value(test_name, &format!("{prefix}[{i}]"), item, out);
            }
        }
        Value::Object(map) => {
            for (key, child) in map {
                scan_value(test_name, &format!("{prefix}.{key}"), child, out);
            }
        }
        _ => {}
    }
}

/// Pull the `(target, value)` of every string exact-value assertion out
/// of a rendered test block by re-parsing it. Re-parsing keeps the scan
/// honest: it sees exactly what landed in YAML, nothing else.
fn exact_assertion_strings(block: &str) -> Vec<(String, String)> {
    let doc = format!("tools:\n{block}");
    let Ok(parsed) = serde_yaml::from_str::<serde_yaml::Value>(&doc) else {
        return Vec::new();
    };
    let Ok(as_json) = serde_json::to_value(parsed) else {
        return Vec::new();
    };
    let Some(entries) = as_json["tools"][0]["expect"]["assertions"].as_array() else {
        return Vec::new();
    };
    entries
        .iter()
        .filter_map(|e| {
            let target = e["target"].as_str()?;
            let exact = e["matcher"]["exact"].as_str()?;
            Some((target.to_string(), exact.to_string()))
        })
        .collect()
}

/// Assemble the final document: schema header, provenance and skip
/// comments, the `servers:` block, and the test blocks under `tools:`.
fn render_suite(opts: &DistillOptions, extraction: &Extraction, blocks: &[String]) -> String {
    let mut out = String::new();
    out.push_str("# yaml-language-server: $schema=https://mcptest.sh/schema/v1.json\n");
    out.push_str(&format!(
        "# Draft suite distilled by `mcptest distill` from {}.\n",
        opts.cassette_path
    ));
    out.push_str(
        "# Args and assertions are lifted verbatim from the recording (after secret\n\
         # redaction). Review every test, especially the `# review:` flagged ones.\n",
    );
    if !extraction.skipped.is_empty() {
        out.push_str(&format!(
            "# Exchanges not distilled: {}\n",
            extraction.skipped.join(", ")
        ));
    }
    out.push('\n');
    out.push_str(&extraction.servers_block);
    out.push('\n');
    out.push_str("tools:\n");
    for (i, block) in blocks.iter().enumerate() {
        if i > 0 {
            out.push('\n');
        }
        out.push_str(block);
    }
    out
}

/// Double-quoted YAML scalar with the escapes a double-quoted scalar
/// needs. Mirrors the propose renderer's quoting so distilled names and
/// paths round-trip through any YAML parser.
pub(crate) fn yaml_quote(s: &str) -> String {
    let mut out = String::with_capacity(s.len() + 2);
    out.push('"');
    for c in s.chars() {
        match c {
            '\\' => out.push_str("\\\\"),
            '"' => out.push_str("\\\""),
            '\n' => out.push_str("\\n"),
            '\r' => out.push_str("\\r"),
            '\t' => out.push_str("\\t"),
            other => out.push(other),
        }
    }
    out.push('"');
    out
}

#[cfg(test)]
mod tests;
