//! Handler for `mcptest ablation`: with/without-tools uplift over a suite.
//!
//! For each single-prompt agent test, this drives the agent loop twice against
//! the same connected server, once with the tools advertised and once with them
//! withheld ([`mcptest_agent::AgentConfig::withhold_server_tools`]), scores each
//! arm against the agent's deterministic `expect:` block, and hands the paired
//! outcomes to the pure [`mcptest_core::ablation`] analyzer. The report says
//! whether attaching the server's tools actually helped.
//!
//! Multi-turn (`turns:`) agents are skipped with a note: ablation pairs a single
//! task across the two arms, and a scripted conversation is a different shape.

use std::sync::Arc;

use anyhow::{anyhow, Context, Result};

use mcptest_agent::stub::StubProvider;
use mcptest_agent::{
    run_agent, AgentConfig, Budget, ConversationTrace, ProviderLookup, ServerHandle,
};
use mcptest_core::ablation::{analyze, render, ArmOutcome, TaskAblation};
use mcptest_core::connector::{connect_server, ServerPool};
use mcptest_core::eval::provider::{LlmProvider, Usage};
use mcptest_core::executor::eval_metric_assertion;
use mcptest_core::network::ProxyConfig;
use mcptest_core::suite::{AgentTest, ResolvedSuite};

use crate::cli::args::AblationArgs;

/// Run the with/without-tools ablation over a suite.
pub(crate) fn ablation_command(args: AblationArgs) -> Result<i32> {
    let runtime = tokio::runtime::Runtime::new().context("could not build the tokio runtime")?;
    runtime.block_on(run(args))
}

async fn run(args: AblationArgs) -> Result<i32> {
    let suite = load_suite(&args.config)?;
    let proxy = ProxyConfig::default();
    let pool = connect_servers(&suite, &proxy).await?;

    let mut tasks = Vec::new();
    for agent in &suite.agents {
        if agent.turns.is_some() {
            eprintln!(
                "ablation: skipping agent `{}` (multi-turn agents are not ablated)",
                agent.name
            );
            continue;
        }
        tasks.push(ablate_agent(agent, &suite, &pool, &proxy, args.allow_stub_provider).await?);
    }

    if tasks.is_empty() {
        eprintln!(
            "ablation: no single-prompt agent tests to ablate in {}",
            args.config.display()
        );
        return Ok(2);
    }

    let report = analyze(&tasks);
    if args.json {
        println!("{}", serde_json::to_string_pretty(&report)?);
    } else {
        print!("{}", render(&report));
    }
    Ok(if report.passed(args.min_uplift) { 0 } else { 1 })
}

/// Read, interpolate, and parse the suite YAML.
fn load_suite(path: &std::path::Path) -> Result<ResolvedSuite> {
    let raw = std::fs::read_to_string(path)
        .with_context(|| format!("could not read {}", path.display()))?;
    let yaml: serde_yaml::Value = serde_yaml::from_str(&raw)
        .with_context(|| format!("could not parse YAML at {}", path.display()))?;
    let (yaml, _secrets) =
        crate::cli::handlers::interp::interpolate_run_config_for_run(yaml, path, true)?;
    mcptest_core::suite::parse(&yaml).map_err(|e| anyhow!("suite parse failed: {e}"))
}

/// Connect every server an agent references into a pool.
async fn connect_servers(suite: &ResolvedSuite, proxy: &ProxyConfig) -> Result<ServerPool> {
    let pool = ServerPool::new();
    let mut seen = std::collections::BTreeSet::new();
    for agent in &suite.agents {
        for name in &agent.servers {
            if !seen.insert(name.clone()) {
                continue;
            }
            let spec = suite
                .servers
                .get(name)
                .with_context(|| format!("agent references unknown server `{name}`"))?;
            let live = connect_server(
                name,
                spec,
                "mcptest",
                mcptest_core::version::semver(),
                proxy,
                &[],
            )
            .await
            .with_context(|| format!("connect to `{name}` failed"))?;
            pool.insert(name.clone(), live).await;
        }
    }
    Ok(pool)
}

/// Whether the provider for an arm is a real (priced) provider or the stub.
enum ProviderSource {
    Real(Arc<dyn LlmProvider>),
    Stub,
}

/// Ablate one single-prompt agent: run both arms and pair their outcomes.
async fn ablate_agent(
    agent: &AgentTest,
    suite: &ResolvedSuite,
    pool: &ServerPool,
    proxy: &ProxyConfig,
    allow_stub: bool,
) -> Result<TaskAblation> {
    let model_ref = agent
        .models
        .first()
        .with_context(|| format!("agent `{}` declares no model", agent.name))?;
    let model_id = model_ref.model_id().to_string();
    let source = resolve_source(agent, suite, proxy, allow_stub)?;

    let resolved = resolve_handles(agent, pool).await?;
    let handles: Vec<ServerHandle<'_>> = resolved
        .iter()
        .map(|(name, srv)| ServerHandle {
            name: name.as_str(),
            client: &srv.client,
        })
        .collect();

    let with_tools = run_arm(agent, &model_id, &source, &handles, false).await?;
    let without_tools = run_arm(agent, &model_id, &source, &handles, true).await?;
    Ok(TaskAblation {
        task: agent.name.clone(),
        with_tools,
        without_tools,
    })
}

/// Resolve the provider source once. A real provider is reused across both arms
/// (it is stateless); the stub is rebuilt per arm (its reply queue drains), so
/// only the source kind is decided here.
fn resolve_source(
    agent: &AgentTest,
    suite: &ResolvedSuite,
    proxy: &ProxyConfig,
    allow_stub: bool,
) -> Result<ProviderSource> {
    let model_ref = agent.models.first().expect("checked by caller");
    match super::run_live_provider::resolve_agent_lookup(model_ref, agent, suite, proxy)? {
        ProviderLookup::Resolved { provider, .. } => Ok(ProviderSource::Real(provider)),
        _ if allow_stub => Ok(ProviderSource::Stub),
        ProviderLookup::MissingCredential { env_var, label } => Err(anyhow!(
            "agent `{}`: no key for {label} (set {env_var}, or pass --allow-stub-provider for an \
             offline wiring smoke test)",
            agent.name
        )),
        ProviderLookup::UnsupportedFamily { .. } => Err(anyhow!(
            "agent `{}`: unsupported model family; use a supported model id, or pass \
             --allow-stub-provider",
            agent.name
        )),
    }
}

/// Resolve an agent's server handles from the pool.
async fn resolve_handles(
    agent: &AgentTest,
    pool: &ServerPool,
) -> Result<Vec<(String, Arc<mcptest_core::connector::LiveServer>)>> {
    let mut out = Vec::with_capacity(agent.servers.len());
    for name in &agent.servers {
        let live = pool.get(name).await.with_context(|| {
            format!("agent `{}` references unknown server `{name}`", agent.name)
        })?;
        out.push((name.clone(), live));
    }
    Ok(out)
}

/// Drive one arm and project the trace onto an ablation outcome.
async fn run_arm(
    agent: &AgentTest,
    model_id: &str,
    source: &ProviderSource,
    handles: &[ServerHandle<'_>],
    withhold: bool,
) -> Result<ArmOutcome> {
    let (provider, real): (Arc<dyn LlmProvider>, bool) = match source {
        ProviderSource::Real(p) => (Arc::clone(p), true),
        ProviderSource::Stub => (
            Arc::new(StubProvider::text(format!(
                "stub model response for agent `{}`",
                agent.name
            ))),
            false,
        ),
    };
    let cfg = AgentConfig {
        model: model_id.to_string(),
        user_prompt: agent.prompt.clone(),
        system_prompt: agent.system_prompt.clone(),
        max_turns: agent.max_turns,
        max_tokens: agent.max_tokens,
        withhold_server_tools: withhold,
        ..AgentConfig::default()
    };
    let budget = Arc::new(Budget::new(None, None));
    let trace = run_agent(&cfg, provider.as_ref(), handles, budget)
        .await
        .map_err(|e| {
            anyhow!(
                "agent `{}` run failed: {}",
                agent.name,
                mcptest_agent::redact_secrets(&e.to_string())
            )
        })?;
    let envelope = trace.to_envelope();
    // An arm passes when every one of the agent's deterministic expectations
    // passes against the trace. An agent with no expectations has no oracle, so
    // neither arm passes and the task reads as unchanged (zero uplift).
    let passed = !agent.expectations.is_empty()
        && agent
            .expectations
            .iter()
            .all(|a| eval_metric_assertion(&envelope, a).is_ok());
    Ok(ArmOutcome {
        passed,
        tool_calls: trace.tool_calls.len(),
        cost_usd: arm_cost(real, provider.as_ref(), &trace, model_id),
        duration_ms: trace.conversation.duration_ms,
    })
}

/// Dollar cost of an arm, from the trace's token counts. Zero for the stub,
/// which does not price a run.
fn arm_cost(real: bool, provider: &dyn LlmProvider, trace: &ConversationTrace, model: &str) -> f64 {
    if !real {
        return 0.0;
    }
    let usage = Usage {
        input_tokens: trace.conversation.tokens.prompt,
        output_tokens: trace.conversation.tokens.completion,
        cache_read_tokens: None,
        cache_write_tokens: None,
    };
    let cost = provider.estimate_cost(&usage, model);
    cost.input_dollars
        + cost.output_dollars
        + cost.cache_read_dollars.unwrap_or(0.0)
        + cost.cache_write_dollars.unwrap_or(0.0)
}
