//! `mcptest generate stubs` standalone scaffolder.
//!
//! Existing tooling builds per-tool YAML stubs as part of `mcptest init`. This
//! module exposes the same generator as a standalone surface so a developer
//! adopting mcptest on an existing repo can scaffold stubs without running the
//! wizard, and so a CI job can regenerate them and diff against checked-in
//! files as a cheap drift check.
//!
//! The generator itself lives in [`mcptest_core::stub_gen`]; this module
//! orchestrates connecting to the server, calling `tools/list`, writing one
//! YAML file per tool, and handling `--overwrite` and `--check` semantics.

use std::collections::BTreeSet;
use std::fs;
use std::path::{Path, PathBuf};
use std::time::Duration;

use anyhow::{Context, Result};
use mcptest_core::protocol::{Client, ProtocolVersion};
use mcptest_core::stub_gen;
use mcptest_core::transport::{StdioConfig, StdioTransport, StreamableHttp, StreamableHttpConfig};
use serde_json::{json, Value};

/// Plain-data arguments for `mcptest generate stubs`.
///
/// Held as its own struct so the CLI dispatcher can fill it in from
/// clap-derived flags and hand the populated struct to [`run_stubs`], and so
/// tests can drive the command without a process boundary.
#[derive(Debug, Clone)]
pub struct GenerateStubsArgs {
    /// MCP server URL to introspect. Exactly one of `url` or
    /// `server_command` must be set.
    pub url: Option<String>,
    /// Argv of a stdio MCP server to spawn and introspect (WOR-1252). Empty
    /// means no stdio target; exactly one of `url` or `server_command` must
    /// be set.
    pub server_command: Vec<String>,
    /// Server key written into each stub's `servers:` block and referenced by
    /// every generated test's `server:` field.
    pub server_name: String,
    /// Optional bearer token forwarded onto every request.
    pub bearer_token: Option<String>,
    /// Directory the generated YAML files are written under.
    pub output_dir: PathBuf,
    /// When `true`, replace any existing stub files. Default is additive:
    /// existing files are left alone so a re-run does not clobber edits.
    pub overwrite: bool,
    /// When `true`, do not write to disk. Emit each stub to stdout
    /// concatenated. Useful for piping into `less` or `git diff -`.
    pub stdout: bool,
    /// When `true`, exit 6 if any generated stub differs from the checked-in
    /// file. Mutually exclusive with `--overwrite`; the CLI rejects that
    /// combination at parse time.
    pub check: bool,
    /// Emit the output-schema conformance test for tools that declare an
    /// `outputSchema`. `--no-output-schema-tests` clears it.
    pub include_output_schema_tests: bool,
    /// Probe tier (`--probe`): also emit deterministic boundary and negative
    /// probe cases per tool, with the documented default seed and cap.
    pub probe: bool,
    /// Suppress per-file progress lines. Honors the global `--quiet` flag.
    pub quiet: bool,
}

/// Synchronous entry point for `mcptest generate stubs <flags>`.
///
/// Connects to the server, calls `tools/list`, runs the stub generator, and
/// writes (or diffs, or prints) one YAML per tool. Returns the desired exit
/// code: 0 on success, 1 on transport or protocol failure, 6 when
/// `--check` finds drift.
pub fn run_stubs(args: GenerateStubsArgs) -> Result<i32> {
    if args.check && args.overwrite {
        eprintln!("mcptest generate stubs: --check and --overwrite are mutually exclusive");
        return Ok(2);
    }
    if args.url.is_none() && args.server_command.is_empty() {
        eprintln!(
            "mcptest generate stubs: a target is required: pass --url <URL> or --server-command \"<cmd>\""
        );
        return Ok(2);
    }
    if args.url.is_some() && !args.server_command.is_empty() {
        eprintln!("mcptest generate stubs: --url and --server-command are mutually exclusive");
        return Ok(2);
    }
    let catalog = match fetch_catalog(&args) {
        Ok(c) => c,
        Err(err) => {
            eprintln!("mcptest generate stubs: {err:#}");
            return Ok(1);
        }
    };
    let opts = stub_gen::StubOptions {
        server_name: args.server_name.clone(),
        include_output_schema_tests: args.include_output_schema_tests,
        probe: args.probe.then(stub_gen::ProbeOptions::default),
    };
    let stubs = match stub_gen::generate_with(&catalog, &opts) {
        Ok(s) => s,
        Err(err) => {
            eprintln!("mcptest generate stubs: catalog had no tools[]: {err}");
            return Ok(1);
        }
    };
    if args.stdout {
        for stub in &stubs {
            println!("# --- {} ---", stub.tool_name);
            print!("{}", stub.yaml);
        }
        return Ok(0);
    }
    if args.check {
        return check_mode(&args, &stubs);
    }
    write_mode(&args, &stubs)
}

fn check_mode(args: &GenerateStubsArgs, stubs: &[stub_gen::ToolStub]) -> Result<i32> {
    let mut drifted: Vec<String> = Vec::new();
    let mut missing: Vec<String> = Vec::new();
    for stub in stubs {
        let path = stub_path(&args.output_dir, &stub.tool_name);
        match fs::read_to_string(&path) {
            Ok(existing) if existing == stub.yaml => {}
            Ok(_) => drifted.push(stub.tool_name.clone()),
            Err(_) => missing.push(stub.tool_name.clone()),
        }
    }
    if drifted.is_empty() && missing.is_empty() {
        if !args.quiet {
            println!(
                "{} stubs match the checked-in files. No drift detected.",
                stubs.len()
            );
        }
        return Ok(0);
    }
    for name in &missing {
        eprintln!("missing stub: {name}");
    }
    for name in &drifted {
        eprintln!("drifted stub: {name}");
    }
    Ok(6)
}

fn write_mode(args: &GenerateStubsArgs, stubs: &[stub_gen::ToolStub]) -> Result<i32> {
    let mut created: Vec<PathBuf> = Vec::new();
    let mut skipped: Vec<PathBuf> = Vec::new();
    let mut overwritten: Vec<PathBuf> = Vec::new();
    fs::create_dir_all(&args.output_dir)
        .with_context(|| format!("failed to create {}", args.output_dir.display()))?;
    for stub in stubs {
        let path = stub_path(&args.output_dir, &stub.tool_name);
        let exists = path.exists();
        if exists && !args.overwrite {
            skipped.push(path);
            continue;
        }
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent)
                .with_context(|| format!("failed to create {}", parent.display()))?;
        }
        fs::write(&path, &stub.yaml)
            .with_context(|| format!("failed to write {}", path.display()))?;
        if exists {
            overwritten.push(path);
        } else {
            created.push(path);
        }
    }
    if !args.quiet {
        for path in &created {
            println!("Generated {}", path.display());
        }
        for path in &overwritten {
            println!("Overwrote {}", path.display());
        }
        for path in &skipped {
            println!("Skipped {} (already exists)", path.display());
        }
        println!(
            "{} stub file(s): {} created, {} overwritten, {} skipped",
            stubs.len(),
            created.len(),
            overwritten.len(),
            skipped.len()
        );
    }
    Ok(0)
}

/// Resolve the on-disk path for a generated stub.
///
/// One YAML per tool lives under `<output_dir>/<safe_name>.yml`. We replace
/// any character outside `[A-Za-z0-9_-]` with `_` so a tool named
/// `tools/list` does not blow up the path traversal.
fn stub_path(output_dir: &Path, tool_name: &str) -> PathBuf {
    let safe: String = tool_name
        .chars()
        .map(|c| match c {
            'a'..='z' | 'A'..='Z' | '0'..='9' | '_' | '-' => c,
            _ => '_',
        })
        .collect();
    output_dir.join(format!("{safe}.yml"))
}

/// Fetch the `tools/list` catalog from whichever target shape the caller
/// set. The two paths share the protocol [`Client`], so the catalog (and the
/// stubs synthesized from it) is identical for the same server either way.
fn fetch_catalog(args: &GenerateStubsArgs) -> Result<Value> {
    match &args.url {
        Some(url) => fetch_catalog_http(url, args.bearer_token.as_deref()),
        None => fetch_catalog_stdio(&args.server_command),
    }
}

fn fetch_catalog_http(url: &str, bearer_token: Option<&str>) -> Result<Value> {
    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .context("failed to build tokio runtime")?;
    runtime.block_on(async {
        let config = StreamableHttpConfig {
            url: url.to_string(),
            bearer_token: bearer_token.map(str::to_string),
            request_timeout: Duration::from_secs(15),
            connect_timeout: Duration::from_secs(15),
            ..StreamableHttpConfig::default()
        };
        let (transport, handle) = StreamableHttp::connect(config)
            .map_err(|e| anyhow::anyhow!("streamable HTTP setup failed: {e}"))?;
        let client = Client::builder(transport)
            .with_request_timeout(Duration::from_secs(15))
            .build();
        client.initialize(ProtocolVersion::current()).await?;
        let response = client.request("tools/list", json!({})).await?;
        client.close().await.ok();
        handle.shutdown().await;
        // The generator expects a `{ tools: [...] }` envelope. Pass through
        // whatever the wire returned: a well-formed server already wraps the
        // array under the same key.
        Ok::<_, anyhow::Error>(response)
    })
}

/// Spawn `command` as a stdio MCP server (the same transport `mcptest run`
/// uses for `command:` servers), fetch `tools/list`, and shut the child
/// down. This is the human-invoked path, so no gating applies; the
/// developer typed the command themselves.
fn fetch_catalog_stdio(command: &[String]) -> Result<Value> {
    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .context("failed to build tokio runtime")?;
    runtime.block_on(async {
        let (program, rest) = command
            .split_first()
            .context("--server-command resolved to an empty argv")?;
        let config = StdioConfig {
            command: program.clone(),
            args: rest.to_vec(),
            ..StdioConfig::default()
        };
        let (transport, handle) = StdioTransport::spawn(config)
            .map_err(|e| anyhow::anyhow!("stdio server setup failed: {e}"))?;
        let client = Client::builder(transport)
            .with_request_timeout(Duration::from_secs(15))
            .build();
        client.initialize(ProtocolVersion::current()).await?;
        let response = client.request("tools/list", json!({})).await?;
        client.close().await.ok();
        // Best-effort teardown: the catalog is already in hand, and a child
        // that ignores EOF is force-killed by the transport.
        let _ = handle.shutdown().await;
        Ok::<_, anyhow::Error>(response)
    })
}

/// Compute the set of tool names that would generate stubs for the given
/// catalog. Exposed for tests so callers can pin behavior without running
/// the full generator pipeline.
pub fn discover_tool_names(catalog: &Value) -> BTreeSet<String> {
    catalog
        .get("tools")
        .and_then(Value::as_array)
        .map(|arr| {
            arr.iter()
                .filter_map(|t| t.get("name").and_then(Value::as_str))
                .map(str::to_string)
                .collect()
        })
        .unwrap_or_default()
}

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

    #[test]
    fn stub_path_sanitizes_unsafe_characters() {
        let p = stub_path(Path::new("/tmp/out"), "tools/list");
        assert_eq!(p.file_name().unwrap(), "tools_list.yml");
    }

    #[test]
    fn discover_tool_names_returns_alphabetical_set() {
        let catalog = json!({
            "tools": [
                { "name": "search" },
                { "name": "ping" },
                { "name": "list_issues" },
            ]
        });
        let names = discover_tool_names(&catalog);
        let collected: Vec<&str> = names.iter().map(String::as_str).collect();
        assert_eq!(collected, vec!["list_issues", "ping", "search"]);
    }

    #[test]
    fn write_mode_creates_one_yaml_per_tool() {
        let dir = tempdir().unwrap();
        let stubs = vec![
            stub_gen::ToolStub {
                tool_name: "alpha".into(),
                yaml: "alpha: yaml\n".into(),
            },
            stub_gen::ToolStub {
                tool_name: "beta".into(),
                yaml: "beta: yaml\n".into(),
            },
        ];
        let args = GenerateStubsArgs {
            url: None,
            server_command: vec![],
            server_name: "server".into(),
            bearer_token: None,
            output_dir: dir.path().to_path_buf(),
            overwrite: false,
            stdout: false,
            check: false,
            include_output_schema_tests: true,
            probe: false,
            quiet: true,
        };
        let code = write_mode(&args, &stubs).unwrap();
        assert_eq!(code, 0);
        assert!(dir.path().join("alpha.yml").exists());
        assert!(dir.path().join("beta.yml").exists());
    }

    #[test]
    fn additive_mode_skips_existing_files() {
        let dir = tempdir().unwrap();
        fs::write(dir.path().join("alpha.yml"), "DO NOT TOUCH").unwrap();
        let stubs = vec![stub_gen::ToolStub {
            tool_name: "alpha".into(),
            yaml: "fresh: yaml\n".into(),
        }];
        let args = GenerateStubsArgs {
            url: None,
            server_command: vec![],
            server_name: "server".into(),
            bearer_token: None,
            output_dir: dir.path().to_path_buf(),
            overwrite: false,
            stdout: false,
            check: false,
            include_output_schema_tests: true,
            probe: false,
            quiet: true,
        };
        write_mode(&args, &stubs).unwrap();
        let body = fs::read_to_string(dir.path().join("alpha.yml")).unwrap();
        assert_eq!(body, "DO NOT TOUCH");
    }

    #[test]
    fn overwrite_mode_replaces_existing_files() {
        let dir = tempdir().unwrap();
        fs::write(dir.path().join("alpha.yml"), "stale").unwrap();
        let stubs = vec![stub_gen::ToolStub {
            tool_name: "alpha".into(),
            yaml: "fresh\n".into(),
        }];
        let args = GenerateStubsArgs {
            url: None,
            server_command: vec![],
            server_name: "server".into(),
            bearer_token: None,
            output_dir: dir.path().to_path_buf(),
            overwrite: true,
            stdout: false,
            check: false,
            include_output_schema_tests: true,
            probe: false,
            quiet: true,
        };
        write_mode(&args, &stubs).unwrap();
        let body = fs::read_to_string(dir.path().join("alpha.yml")).unwrap();
        assert_eq!(body, "fresh\n");
    }

    #[test]
    fn check_mode_returns_zero_when_files_match() {
        let dir = tempdir().unwrap();
        fs::write(dir.path().join("alpha.yml"), "ok\n").unwrap();
        let stubs = vec![stub_gen::ToolStub {
            tool_name: "alpha".into(),
            yaml: "ok\n".into(),
        }];
        let args = GenerateStubsArgs {
            url: None,
            server_command: vec![],
            server_name: "server".into(),
            bearer_token: None,
            output_dir: dir.path().to_path_buf(),
            overwrite: false,
            stdout: false,
            check: true,
            include_output_schema_tests: true,
            probe: false,
            quiet: true,
        };
        let code = check_mode(&args, &stubs).unwrap();
        assert_eq!(code, 0);
    }

    #[test]
    fn check_mode_returns_six_when_file_drifts() {
        let dir = tempdir().unwrap();
        fs::write(dir.path().join("alpha.yml"), "old\n").unwrap();
        let stubs = vec![stub_gen::ToolStub {
            tool_name: "alpha".into(),
            yaml: "new\n".into(),
        }];
        let args = GenerateStubsArgs {
            url: None,
            server_command: vec![],
            server_name: "server".into(),
            bearer_token: None,
            output_dir: dir.path().to_path_buf(),
            overwrite: false,
            stdout: false,
            check: true,
            include_output_schema_tests: true,
            probe: false,
            quiet: true,
        };
        let code = check_mode(&args, &stubs).unwrap();
        assert_eq!(code, 6);
    }
}
