//! Handler for `mcptest security`.

use std::path::Path;

use anyhow::Result;
use mcptest_core::cassette::portable::Cassette;
use mcptest_core::security::skill_scan::{parse_skill, scan_skill};
use mcptest_core::security::{
    normalize_generic_json, normalize_sarif, normalize_snyk, probe_evidence_from_snapshot,
    render_run_json, render_run_pretty_styled, render_sarif, run_cassette, run_deterministic,
    run_import, run_with_advisory, servers_from_snapshot, SecurityEngine, SecurityFinding,
    SecurityRun, SecurityRunOptions, SupplementFinding,
};
use mcptest_core::style::Styler;
use mcptest_core::waivers::{partition, render_summary, CivilDate, WaiverFile};
use serde_json::Value;

use crate::cli::args::{
    SecurityArgs, SecurityCassetteArgs, SecurityCommand, SecurityFormat, SecurityImportArgs,
    SecuritySkillArgs,
};
use crate::cli::global::GlobalArgs;

mod adaptive;
mod corpus;
mod judge;
mod proof;
mod redteam;

/// Handle `mcptest security snapshot.json` and `mcptest security redteam`.
///
/// The snapshot form loads a `tools/list`-style JSON snapshot and runs the
/// deterministic lanes (surface, toxic-flow, plus namespace when the snapshot
/// carries a `servers` array and integrity when `--baseline` is given). The
/// deterministic findings own the verdict (ADR 0040). With `--model`, the
/// advisory LLM-judge lane also runs, reported in its own section and never
/// affecting the verdict (ADR 0041). The `redteam` subcommand routes to the
/// live red-team lane, whose exploitability signal is advisory only (ADR 0041).
///
/// Exit code 2 when an input cannot be read or parsed, 1 when any deterministic
/// finding is at or above `--fail-on`, otherwise 0.
pub(crate) fn security_command(args: SecurityArgs, global: &GlobalArgs) -> Result<i32> {
    match args.command {
        Some(SecurityCommand::Redteam(redteam_args)) => {
            return redteam::redteam_command(redteam_args, global);
        }
        Some(SecurityCommand::Import(import_args)) => {
            return import_command(import_args, global);
        }
        Some(SecurityCommand::Cassette(cassette_args)) => {
            return cassette_command(cassette_args, global);
        }
        Some(SecurityCommand::Proof(proof_args)) => {
            return proof::proof_command(proof_args);
        }
        Some(SecurityCommand::Skill(skill_args)) => {
            return skill_command(skill_args);
        }
        None => {}
    }
    // `subcommand_negates_reqs` makes the positional optional in the type, but
    // clap requires it whenever no subcommand is given, so this is always set
    // on the snapshot path.
    let Some(snapshot_path) = &args.snapshot else {
        eprintln!("mcptest security: a snapshot path or a subcommand is required");
        return Ok(2);
    };
    let snapshot = match load_json(snapshot_path) {
        Ok(v) => v,
        Err(code) => return Ok(code),
    };

    let mut options = SecurityRunOptions {
        servers: servers_from_snapshot(&snapshot),
        probe_evidence: probe_evidence_from_snapshot(&snapshot),
        ..Default::default()
    };
    if let Some(baseline_path) = &args.baseline {
        match load_json(baseline_path) {
            Ok(v) => options.baseline = Some(v),
            Err(code) => return Ok(code),
        }
    }

    let run = match &args.model {
        Some(model) => match judge::build_judge(model) {
            Ok(judge) => run_with_advisory(&snapshot, &options, judge),
            Err(e) => {
                eprintln!("mcptest security: cannot enable the advisory lane: {e}");
                return Ok(2);
            }
        },
        None => run_deterministic(&snapshot, &options),
    };

    render(
        args.format,
        &run,
        global
            .color
            .resolve(std::io::IsTerminal::is_terminal(&std::io::stdout())),
    )?;

    let threshold = args.fail_on.to_severity();
    let findings_fail = match &args.waivers {
        Some(path) => {
            let body = match std::fs::read_to_string(path) {
                Ok(b) => b,
                Err(e) => {
                    eprintln!("mcptest security: cannot read waiver file {path:?}: {e}");
                    return Ok(2);
                }
            };
            let file = match WaiverFile::parse(&body) {
                Ok(f) => f,
                Err(e) => {
                    eprintln!("mcptest security: {path:?} is not a valid waiver file: {e}");
                    return Ok(2);
                }
            };
            let (active, waived) =
                partition(run.report.findings.clone(), &file, CivilDate::today_utc());
            let summary = render_summary(&waived);
            if !summary.is_empty() {
                eprint!("{summary}");
            }
            active.iter().any(|f| f.severity >= threshold)
        }
        None => run.report.findings.iter().any(|f| f.severity >= threshold),
    };
    // A supplement-only failure (an imported scanner finding) still gates even
    // when waivers clear the deterministic findings; derive it from the public
    // gate without re-deriving the supplement rule.
    let supplement_fail =
        run.fails_at(threshold) && !run.report.findings.iter().any(|f| f.severity >= threshold);
    if findings_fail || supplement_fail {
        Ok(1)
    } else {
        Ok(0)
    }
}

/// Handle `mcptest security skill`: scan agent skill files for instruction-text
/// attacks. Exit 2 on a read error, 1 when a finding meets `--fail-on`, else 0.
fn skill_command(args: SecuritySkillArgs) -> Result<i32> {
    let mut all: Vec<SecurityFinding> = Vec::new();
    for path in &args.paths {
        let text = match std::fs::read_to_string(path) {
            Ok(t) => t,
            Err(e) => {
                eprintln!("mcptest security skill: cannot read {path:?}: {e}");
                return Ok(2);
            }
        };
        let doc = parse_skill(&text);
        all.extend(scan_skill(&doc.name, &doc.body));
    }

    if args.json {
        println!("{}", serde_json::to_string_pretty(&all)?);
    } else if all.is_empty() {
        println!("no agent-skill findings");
    } else {
        for f in &all {
            println!(
                "{:<8} {:<8} {} :: {}",
                f.severity.as_str(),
                f.rule_id,
                f.definition,
                f.message
            );
            if !f.evidence.is_empty() {
                println!("           evidence: {}", f.evidence);
            }
        }
    }

    if let Some(fail_on) = args.fail_on {
        let threshold = fail_on.to_severity();
        if all.iter().any(|f| f.severity >= threshold) {
            return Ok(1);
        }
    }
    Ok(0)
}

/// Handle `mcptest security import`.
///
/// Each scanner file is normalized into the finding vocabulary: SARIF via
/// [`normalize_sarif`], Snyk agent-scan JSON via [`normalize_snyk`], and any
/// other shape via [`normalize_generic_json`]. When `--snapshot` is given, the
/// bundled deterministic lanes run too and the imports are folded in beside
/// them; otherwise the imports stand alone. [`run_import`] dedups every import
/// against the bundled catalog, so an overlap is counted once.
///
/// Exit code 2 when a file cannot be read or parsed or when no scanner file is
/// given, 1 when any counted finding is at or above `--fail-on`, otherwise 0.
fn import_command(args: SecurityImportArgs, global: &GlobalArgs) -> Result<i32> {
    if args.sarif.is_empty() && args.snyk.is_empty() && args.supplement.is_empty() {
        eprintln!(
            "mcptest security import: at least one of --sarif, --snyk, or --supplement is required"
        );
        return Ok(2);
    }

    let mut supplements: Vec<SupplementFinding> = Vec::new();
    for path in &args.sarif {
        let doc = match load_json(path) {
            Ok(v) => v,
            Err(code) => return Ok(code),
        };
        let tool = sarif_tool_name(&doc);
        supplements.extend(normalize_sarif(&doc, &tool, args.advisory));
    }
    for path in &args.snyk {
        let doc = match load_json(path) {
            Ok(v) => v,
            Err(code) => return Ok(code),
        };
        supplements.extend(normalize_snyk(&doc, "snyk", args.advisory));
    }
    for path in &args.supplement {
        let doc = match load_json(path) {
            Ok(v) => v,
            Err(code) => return Ok(code),
        };
        let tool = path
            .file_stem()
            .and_then(|s| s.to_str())
            .unwrap_or("external");
        supplements.extend(normalize_generic_json(&doc, tool, args.advisory));
    }

    let mut options = SecurityRunOptions::default();
    let snapshot = match &args.snapshot {
        Some(path) => match load_json(path) {
            Ok(v) => {
                options.servers = servers_from_snapshot(&v);
                options.probe_evidence = probe_evidence_from_snapshot(&v);
                Some(v)
            }
            Err(code) => return Ok(code),
        },
        None => None,
    };

    let run = run_import(snapshot.as_ref(), &options, supplements);
    render(
        args.format,
        &run,
        global
            .color
            .resolve(std::io::IsTerminal::is_terminal(&std::io::stdout())),
    )?;

    if run.fails_at(args.fail_on.to_severity()) {
        Ok(1)
    } else {
        Ok(0)
    }
}

/// Handle `mcptest security cassette`.
///
/// Loads each recorded cassette and scans its tool-call results and error
/// bodies for output poisoning (SEC-063 injection, SEC-064 hidden unicode),
/// reusing the surface detectors. The findings render and gate through the same
/// path as the snapshot scan.
///
/// Exit code 2 when a cassette cannot be read or parsed, 1 when any finding is
/// at or above `--fail-on`, otherwise 0.
fn cassette_command(args: SecurityCassetteArgs, global: &GlobalArgs) -> Result<i32> {
    let mut cassettes = Vec::with_capacity(args.cassettes.len());
    for path in &args.cassettes {
        let body = match std::fs::read_to_string(path) {
            Ok(b) => b,
            Err(e) => {
                eprintln!("mcptest security: cannot read {path:?}: {e}");
                return Ok(2);
            }
        };
        match serde_json::from_str::<Cassette>(&body) {
            Ok(cassette) => cassettes.push(cassette),
            Err(e) => {
                eprintln!("mcptest security: {path:?} is not a valid cassette: {e}");
                return Ok(2);
            }
        }
    }

    let run = run_cassette(&cassettes);
    render(
        args.format,
        &run,
        global
            .color
            .resolve(std::io::IsTerminal::is_terminal(&std::io::stdout())),
    )?;

    if run.fails_at(args.fail_on.to_severity()) {
        Ok(1)
    } else {
        Ok(0)
    }
}

/// Read a SARIF log's scanner name from `runs[0].tool.driver.name` for
/// provenance, falling back to `external` when the log omits it.
fn sarif_tool_name(sarif: &Value) -> String {
    sarif
        .get("runs")
        .and_then(|v| v.as_array())
        .and_then(|runs| runs.first())
        .and_then(|run| run.get("tool"))
        .and_then(|tool| tool.get("driver"))
        .and_then(|driver| driver.get("name"))
        .and_then(|v| v.as_str())
        .unwrap_or("external")
        .to_string()
}

/// Read and parse a JSON file, mapping read and parse failures to exit code 2.
/// Returns the parsed value or the exit code to return.
fn load_json(path: &Path) -> std::result::Result<Value, i32> {
    let body = match std::fs::read_to_string(path) {
        Ok(b) => b,
        Err(e) => {
            eprintln!("mcptest security: cannot read {path:?}: {e}");
            return Err(2);
        }
    };
    match serde_json::from_str(&body) {
        Ok(v) => Ok(v),
        Err(e) => {
            eprintln!("mcptest security: {path:?} is not valid JSON: {e}");
            Err(2)
        }
    }
}

/// Render the run in the requested format. SARIF is verdict-only: it is for
/// code scanning, where only the deterministic findings belong, so the posture
/// map and advisory section are omitted from it.
fn render(format: SecurityFormat, run: &SecurityRun, color: bool) -> Result<()> {
    match format {
        SecurityFormat::Pretty => print!("{}", render_run_pretty_styled(run, &Styler::new(color))),
        SecurityFormat::Json => println!("{}", render_run_json(run)?),
        SecurityFormat::Sarif => {
            let engine = SecurityEngine::with_bundled_checks();
            println!("{}", render_sarif(engine.checks(), &run.report)?);
        }
        SecurityFormat::Html => {
            let engine = SecurityEngine::with_bundled_checks();
            print!(
                "{}",
                mcptest_core::security::findings_report::render_html(run, engine.checks())
            );
        }
        SecurityFormat::Md => {
            let engine = SecurityEngine::with_bundled_checks();
            print!(
                "{}",
                mcptest_core::security::findings_report::render_markdown(run, engine.checks())
            );
        }
    }
    Ok(())
}
