//! Live `CaptureRunner` for `mcptest model-compat capture` / `run`.
//!
//! The capture engine in `mcptest_core::model_compat` is transport-free: it
//! drives a [`CaptureRunner`] to enumerate and execute tests, then writes a
//! canonical baseline. Until WOR-1385 the only implementor was a test fake, so
//! `capture` and `run` could only ride on a pre-built `--input` / `--candidate`
//! baseline. This module supplies the real implementor.
//!
//! It lives in the binary crate, not in `mcptest-core`, on purpose:
//! `mcptest-agent` depends on `mcptest-core`, so `mcptest-core` cannot depend on
//! `mcptest-agent` without a dependency cycle. The agent driver
//! [`mcptest_agent::run_agent`] is the same primitive the live `agents:` runner
//! and `eval_live::run_live_candidate` use, so a live
//! capture grades a candidate of the identical shape; we reuse the driver
//! rather than reimplement a model loop. We drive [`run_agent`] directly (not
//! the `Value`-returning `run_live_candidate` bridge) because the capture model
//! needs the structured [`ConversationTrace`] (finish reason, per-call args),
//! which the JSON envelope flattens.

use std::sync::Arc;

use async_trait::async_trait;
use mcptest_agent::{run_agent, AgentConfig, Budget, ConversationTrace, ServerHandle};
use mcptest_core::eval::provider::LlmProvider;
use mcptest_core::model_compat::{
    CaptureRunner, CapturedExchange, ContentBlock, FinishReason, ModelCompatError, ToolCall,
};

/// One prompt scenario the live runner captures.
///
/// `id` is the stable assertion id (`<file>::<test>::<step>`) the baseline
/// keys on; `prompt` is the user turn the agent is driven against.
#[derive(Debug, Clone)]
pub struct CaptureScenario {
    /// Stable assertion id stored on the captured exchange.
    pub id: String,
    /// User prompt that kicks off the agent run for this scenario.
    pub prompt: String,
}

/// A [`CaptureRunner`] backed by a live MCP server and a real (or stub) model.
///
/// Each `run_test` drives one scenario's prompt through the agent driver
/// against the connected servers and maps the resulting [`ConversationTrace`]
/// into a [`CapturedExchange`]. The runner borrows its servers and provider so
/// the caller owns connection lifecycle (the server pool tears the wire down at
/// end of run, exactly as the live eval path does).
pub struct LiveCaptureRunner<'a> {
    scenarios: Vec<CaptureScenario>,
    servers: &'a [ServerHandle<'a>],
    provider: &'a dyn LlmProvider,
    model: String,
    budget: Arc<Budget>,
}

impl<'a> LiveCaptureRunner<'a> {
    /// Build a live capture runner over a fixed scenario list.
    ///
    /// `model` is the provider-qualified id passed to the driver; `budget` is
    /// the shared stop-loss the agent loop charges against.
    pub fn new(
        scenarios: Vec<CaptureScenario>,
        servers: &'a [ServerHandle<'a>],
        provider: &'a dyn LlmProvider,
        model: impl Into<String>,
        budget: Arc<Budget>,
    ) -> Self {
        Self {
            scenarios,
            servers,
            provider,
            model: model.into(),
            budget,
        }
    }
}

#[async_trait]
impl CaptureRunner for LiveCaptureRunner<'_> {
    async fn list_tests(&self) -> Result<Vec<String>, ModelCompatError> {
        Ok(self.scenarios.iter().map(|s| s.id.clone()).collect())
    }

    async fn run_test(&self, test_id: &str) -> Result<CapturedExchange, ModelCompatError> {
        let scenario = self
            .scenarios
            .iter()
            .find(|s| s.id == test_id)
            .ok_or_else(|| ModelCompatError::InvalidFilter {
                pattern: test_id.into(),
                reason: "unknown capture scenario".into(),
            })?;
        let config = AgentConfig {
            model: self.model.clone(),
            user_prompt: scenario.prompt.clone(),
            ..AgentConfig::default()
        };
        let trace = run_agent(
            &config,
            self.provider,
            self.servers,
            Arc::clone(&self.budget),
        )
        .await
        .map_err(|e| capture_failed(test_id, &e.to_string()))?;
        Ok(trace_to_exchange(test_id, &trace))
    }
}

/// Wrap an agent-run failure as a [`ModelCompatError`], redacting any secret
/// that leaked into the driver error before it reaches the user.
fn capture_failed(test_id: &str, message: &str) -> ModelCompatError {
    ModelCompatError::InvalidFilter {
        pattern: test_id.into(),
        reason: format!(
            "live capture failed: {}",
            mcptest_agent::redact_secrets(message)
        ),
    }
}

/// Map a finished [`ConversationTrace`] onto the capture model.
///
/// The final text becomes a single text content block (skipped when empty so a
/// pure tool-call run does not carry a blank block); every tool call carries
/// across in order; the finish reason is derived from the run's stop reason and
/// the truncation flag so a run cut at `max_tokens` records `MaxTokens` rather
/// than a clean `Stop`.
fn trace_to_exchange(test_id: &str, trace: &ConversationTrace) -> CapturedExchange {
    let mut content = Vec::new();
    if !trace.final_response.is_empty() {
        content.push(ContentBlock::Text {
            text: trace.final_response.clone(),
        });
    }
    let tool_calls = trace
        .tool_calls
        .iter()
        .map(|c| ToolCall {
            name: c.name.clone(),
            arguments: c.args.clone(),
            raw_arguments: None,
        })
        .collect();
    CapturedExchange {
        id: test_id.to_string(),
        content,
        tool_calls,
        finish_reason: finish_reason_for(trace),
        mcp_messages: Vec::new(),
    }
}

/// Derive the capture finish reason from the trace. A truncated run maps to
/// [`FinishReason::MaxTokens`]; any tool call without a clean text finish maps
/// to [`FinishReason::ToolUse`]; otherwise the run stopped naturally.
fn finish_reason_for(trace: &ConversationTrace) -> FinishReason {
    if trace.truncated() {
        return FinishReason::MaxTokens;
    }
    if trace.final_response.is_empty() && !trace.tool_calls.is_empty() {
        return FinishReason::ToolUse;
    }
    FinishReason::Stop
}

/// Map a suite's agent tests to capture scenarios: each agent's name is the
/// stable assertion id and its user prompt is the turn the model is driven
/// against. This is the live-capture scenario source, so `model-compat
/// capture` records exactly the agent runs the suite already declares.
pub fn scenarios_from_agents(agents: &[mcptest_core::suite::AgentTest]) -> Vec<CaptureScenario> {
    agents
        .iter()
        .map(|a| CaptureScenario {
            id: a.name.clone(),
            prompt: a.prompt.clone(),
        })
        .collect()
}

/// Capture a model-compat baseline live (WOR-1385): parse the suite, connect
/// the servers its agents use, resolve a provider for `model` from the
/// environment, drive each agent prompt through the live runner, and write the
/// canonical baseline to `output`. Returns the process exit code. With no
/// provider key resolved, or a suite with no agents, it prints a clear message
/// and returns a non-zero code rather than writing an empty baseline.
pub fn run_live_capture(
    config_path: &std::path::Path,
    model: &str,
    output: &std::path::Path,
    filter: Option<&str>,
    quiet: bool,
) -> anyhow::Result<i32> {
    use anyhow::Context as _;
    let body = std::fs::read_to_string(config_path)
        .with_context(|| format!("could not read {}", config_path.display()))?;
    let yaml: serde_yaml::Value = serde_yaml::from_str(&body)
        .with_context(|| format!("could not parse YAML at {}", config_path.display()))?;
    let yaml =
        crate::cli::handlers::interp::interpolate_run_config_reported(yaml, config_path, quiet)?;
    let suite = mcptest_core::suite::parse(&yaml)
        .map_err(|e| anyhow::anyhow!("typed suite parse failed: {e}"))?;

    let scenarios = scenarios_from_agents(&suite.agents);
    if scenarios.is_empty() {
        eprintln!(
            "mcptest model-compat capture: no agents to capture in {}; a live capture needs an `agents:` block",
            config_path.display()
        );
        return Ok(1);
    }

    let proxy = mcptest_core::network::ProxyConfig::default();
    let provider = match mcptest_agent::from_env(model, &proxy) {
        Ok(mcptest_agent::ProviderLookup::Resolved { provider, .. }) => provider,
        _ => {
            eprintln!(
                "mcptest model-compat capture: no provider key resolved for `{model}`; live capture needs a model API key in the environment"
            );
            return Ok(1);
        }
    };

    let rt = tokio::runtime::Runtime::new().context("could not build tokio runtime")?;
    rt.block_on(async {
        // Connect every server any agent references, once each, into a pool.
        let pool = mcptest_core::connector::ServerPool::new();
        let mut names: Vec<String> = Vec::new();
        for agent in &suite.agents {
            for name in &agent.servers {
                if names.iter().any(|n| n == name) {
                    continue;
                }
                let Some(spec) = suite.servers.get(name) else {
                    return Err(anyhow::anyhow!(
                        "agent `{}` references unknown server `{name}`",
                        agent.name
                    ));
                };
                let live = mcptest_core::connector::connect_server(
                    name,
                    spec,
                    "mcptest",
                    mcptest_core::version::semver(),
                    &proxy,
                    &[],
                )
                .await
                .with_context(|| format!("could not connect server `{name}`"))?;
                pool.insert(name.clone(), live).await;
                names.push(name.clone());
            }
        }
        // Hold the Arcs so the borrowed `ServerHandle`s stay valid for the run.
        let mut servers = Vec::with_capacity(names.len());
        for name in &names {
            if let Some(srv) = pool.get(name).await {
                servers.push((name.clone(), srv));
            }
        }
        let handles: Vec<ServerHandle> = servers
            .iter()
            .map(|(name, srv)| ServerHandle {
                name: name.as_str(),
                client: &srv.client,
            })
            .collect();

        let budget = std::sync::Arc::new(Budget::new(None, None));
        let runner = LiveCaptureRunner::new(scenarios, &handles, provider.as_ref(), model, budget);
        let options = mcptest_core::model_compat::CaptureOptions {
            model_id: model.to_string(),
            filter: filter.map(str::to_string),
            ..Default::default()
        };
        let baseline = mcptest_core::model_compat::capture_baseline(&runner, &options, output)
            .await
            .map_err(|e| anyhow::anyhow!("live capture failed: {e}"))?;
        if !quiet {
            println!(
                "mcptest model-compat capture: wrote {} assertions to {}",
                baseline.assertions.len(),
                output.display()
            );
        }
        Ok(0)
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    // Only the unix-gated live-capture stdio test uses these; gating the
    // imports too keeps the Windows build free of unused-import warnings.
    #[cfg(unix)]
    use mcptest_agent::stub::{StubProvider, StubReply};
    #[cfg(unix)]
    use mcptest_core::model_compat::{capture_baseline, CaptureOptions};
    use serde_json::json;

    #[test]
    fn scenarios_are_sourced_from_agent_name_and_prompt() {
        let yaml: serde_yaml::Value = serde_yaml::from_str(
            "servers:\n  s:\n    command: [\"true\"]\nagents:\n  - name: greet\n    model: stub\n    servers: [s]\n    prompt: say hi\n",
        )
        .expect("yaml parses");
        let suite = mcptest_core::suite::parse(&yaml).expect("suite parses");
        let scenarios = scenarios_from_agents(&suite.agents);
        assert_eq!(scenarios.len(), 1);
        assert_eq!(scenarios[0].id, "greet");
        assert_eq!(scenarios[0].prompt, "say hi");
    }

    // Builtins-only stdio fixture (no forks) so the handshake stays robust
    // under the gate's parallel load, mirroring the eval_live fixture. The
    // server answers initialize, a one-tool tools/list, and any tools/call.
    #[cfg(unix)]
    const FIXTURE_SERVER: &str = r##"#!/bin/sh
while IFS= read -r line; do
  case "$line" in
    *'"id":'*) tmp=${line#*'"id":'}; id=${tmp%%[!0-9]*} ;;
    *) id= ;;
  esac
  case "$line" in
    *'"method":"initialize"'*)
      printf '{"jsonrpc":"2.0","id":%s,"result":{"protocolVersion":"2025-03-26","serverInfo":{"name":"fix","version":"0.0.1"},"capabilities":{}}}\n' "$id" ;;
    *'"method":"notifications/initialized"'*) : ;;
    *'"method":"tools/list"'*)
      printf '{"jsonrpc":"2.0","id":%s,"result":{"tools":[{"name":"echo","description":"echo","inputSchema":{"type":"object"}}]}}\n' "$id" ;;
    *'"method":"tools/call"'*)
      printf '{"jsonrpc":"2.0","id":%s,"result":{"content":[{"type":"text","text":"ok"}]}}\n' "$id" ;;
    *)
      [ -n "$id" ] && printf '{"jsonrpc":"2.0","id":%s,"error":{"code":-32601,"message":"method not found"}}\n' "$id" ;;
  esac
done
"##;

    #[cfg(unix)]
    #[tokio::test]
    async fn live_runner_captures_a_two_scenario_baseline_offline() {
        use mcptest_core::protocol::{Client, ProtocolVersion};
        use mcptest_core::transport::{StdioConfig, StdioTransport};
        use std::os::unix::fs::PermissionsExt;
        use std::time::Duration;

        let tmp = tempfile::TempDir::new().unwrap();
        let script = tmp.path().join("server.sh");
        std::fs::write(&script, FIXTURE_SERVER).unwrap();
        let mut perms = std::fs::metadata(&script).unwrap().permissions();
        perms.set_mode(0o755);
        std::fs::set_permissions(&script, perms).unwrap();

        let cfg = StdioConfig {
            command: script.to_string_lossy().into(),
            args: vec![],
            env: Default::default(),
            working_dir: None,
            shutdown_timeout: Duration::from_secs(2),
        };
        let (transport, handle) = StdioTransport::spawn(cfg).expect("spawn fixture");
        let client = Client::builder(transport).build();
        client
            .initialize(ProtocolVersion::V2025_03_26)
            .await
            .expect("initialize");
        // Leak the child handle for the test's duration: dropping it would kill
        // the server before the runner finishes calling it.
        Box::leak(Box::new(handle));

        let handles = [ServerHandle {
            name: "fix",
            client: &client,
        }];
        // Scenario one: the model calls a tool, then answers. Scenario two: a
        // plain text answer. The stub scripts both deterministically so the
        // capture runs with no network and no key.
        let provider = StubProvider::new(vec![
            StubReply::ToolUse {
                name: "echo".into(),
                args: json!({"msg": "hi"}),
            },
            StubReply::Text("done".into()),
            StubReply::Text("just text".into()),
        ]);
        let budget = Arc::new(Budget::new(None, None));
        let runner = LiveCaptureRunner::new(
            vec![
                CaptureScenario {
                    id: "suite::tool::1".into(),
                    prompt: "use the echo tool".into(),
                },
                CaptureScenario {
                    id: "suite::text::1".into(),
                    prompt: "say something".into(),
                },
            ],
            &handles,
            &provider,
            "stub",
            budget,
        );

        let out = tmp.path().join("baseline.json");
        let baseline = capture_baseline(&runner, &CaptureOptions::default(), &out)
            .await
            .expect("capture succeeds offline");

        assert_eq!(baseline.assertions.len(), 2);
        // Sorted by id: text scenario sorts before tool scenario.
        let text = &baseline.assertions[0];
        assert_eq!(text.id, "suite::text::1");
        assert_eq!(text.finish_reason, FinishReason::Stop);
        assert!(text.tool_calls.is_empty());

        let tool = &baseline.assertions[1];
        assert_eq!(tool.id, "suite::tool::1");
        assert_eq!(tool.tool_calls.len(), 1);
        assert_eq!(tool.tool_calls[0].name, "echo");
        assert!(out.exists());
    }

    fn trace_with(
        final_response: &str,
        calls: Vec<mcptest_agent::AgentToolCall>,
    ) -> ConversationTrace {
        use mcptest_agent::trace::ConversationSummary;
        use mcptest_agent::{AgentStopReason, ConversationTokens, ConversationTurn};
        ConversationTrace {
            tool_calls: calls,
            tool_results: vec![],
            final_response: final_response.into(),
            conversation: ConversationSummary {
                tokens: ConversationTokens::default(),
                duration_ms: 0,
                message_count: 0,
                stop_reason: AgentStopReason::Completed,
                turns: Vec::<ConversationTurn>::new(),
                per_turn: vec![],
            },
        }
    }

    #[test]
    fn maps_final_text_to_a_single_content_block() {
        let trace = trace_with("hello", vec![]);
        let ex = trace_to_exchange("t::1", &trace);
        assert_eq!(ex.content.len(), 1);
        assert_eq!(ex.finish_reason, FinishReason::Stop);
    }

    #[test]
    fn empty_final_text_with_a_tool_call_is_tool_use() {
        let call = mcptest_agent::AgentToolCall {
            name: "search".into(),
            server: "s".into(),
            args: json!({"q": "x"}),
            ..Default::default()
        };
        let trace = trace_with("", vec![call]);
        let ex = trace_to_exchange("t::1", &trace);
        assert!(ex.content.is_empty(), "no blank block for an empty answer");
        assert_eq!(ex.tool_calls.len(), 1);
        assert_eq!(ex.finish_reason, FinishReason::ToolUse);
    }
}
