//! Arguments for `mcptest pipe` and its budget controls.

use std::path::PathBuf;
use std::time::Duration;

use clap::{Args, ValueEnum};

use crate::cli::parsers::{parse_dollar_amount, parse_duration, parse_key_value};

/// Arguments for `mcptest pipe <pipeline.yaml>`.
///
/// Runs a declarative multi-step tool-call pipeline. The pipeline file is the
/// one positional argument; `--var` injects values referenced as `${var.X}`;
/// the budget flags cap cumulative spend across every step.
#[derive(Debug, Args)]
pub struct PipeArgs {
    /// Path to the pipeline YAML file.
    #[arg(value_name = "PIPELINE")]
    pub pipeline: PathBuf,

    /// MCP server endpoint the pipeline's tool calls target.
    #[arg(long, value_name = "URL")]
    pub url: Option<String>,

    /// Env var that holds a bearer token to attach to every request.
    #[arg(long = "bearer-token-env", value_name = "NAME")]
    pub bearer_token_env: Option<String>,

    /// Inject a pipeline variable. Repeatable. Format `KEY=VALUE`. Referenced
    /// in the pipeline as `${var.KEY}`.
    #[arg(long = "var", value_name = "KEY=VALUE", value_parser = parse_key_value)]
    pub vars: Vec<(String, String)>,

    /// Output format. `pretty` prints the last step's result; `json` prints
    /// the full execution trace.
    #[arg(long, value_enum, value_name = "FORMAT", default_value_t = PipeFormat::Pretty)]
    pub format: PipeFormat,

    /// Print the planned execution without making any tool calls.
    #[arg(long = "dry-run")]
    pub dry_run: bool,

    /// With `--dry-run`, also print a projected cost estimate.
    #[arg(long = "estimate-cost", requires = "dry_run")]
    pub estimate_cost: bool,

    /// Aggregate USD ceiling across all steps.
    #[arg(long = "max-cost", value_name = "USD", value_parser = parse_dollar_amount)]
    pub max_cost: Option<f64>,

    /// Aggregate token ceiling (input + output) across all steps.
    #[arg(long = "max-tokens", value_name = "N")]
    pub max_tokens: Option<u64>,

    /// Per-step USD ceiling, layered on top of the aggregate cap.
    #[arg(long = "max-cost-per-call", value_name = "USD", value_parser = parse_dollar_amount)]
    pub max_cost_per_call: Option<f64>,

    /// Wall-clock ceiling for the whole run. Accepts `30s`, `2m`,
    /// `1h`.
    #[arg(long = "max-duration", value_name = "DURATION", value_parser = parse_duration)]
    pub max_duration: Option<Duration>,

    /// How to react to a budget breach.
    #[arg(
        long = "on-budget-exceeded",
        value_enum,
        value_name = "MODE",
        default_value_t = OnBudgetExceededCli::Stop
    )]
    pub on_budget_exceeded: OnBudgetExceededCli,

    /// Path to a pricing-table override. Defaults to the bundled
    /// `pricing.yaml`.
    #[arg(long = "pricing-table", value_name = "PATH")]
    pub pricing_table: Option<PathBuf>,
}

/// Output formats for `mcptest pipe`.
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, ValueEnum)]
#[value(rename_all = "kebab-case")]
pub enum PipeFormat {
    /// Print the last step's result.
    #[default]
    Pretty,
    /// Print the full execution trace as JSON.
    Json,
}

/// Budget-breach behavior selector for `--on-budget-exceeded`.
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, ValueEnum)]
#[value(rename_all = "kebab-case")]
pub enum OnBudgetExceededCli {
    /// Stop before the breaching step and exit non-zero (default).
    #[default]
    Stop,
    /// Run to completion, then exit non-zero if a cap was exceeded.
    Continue,
    /// Run to completion, warn on stderr, exit zero.
    Warn,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::cli::{Cli, Command};
    use clap::Parser;

    #[test]
    fn parses_pipe_with_var_and_budget() {
        let cli = Cli::try_parse_from([
            "mcptest",
            "pipe",
            "flow.yaml",
            "--var",
            "Q=alice",
            "--max-cost",
            "$0.50",
        ])
        .expect("parse");
        match cli.command {
            Command::Pipe(args) => {
                assert_eq!(args.pipeline.to_str().unwrap(), "flow.yaml");
                assert_eq!(args.vars, vec![("Q".to_string(), "alice".to_string())]);
                assert_eq!(args.max_cost, Some(0.50));
            }
            other => panic!("expected Pipe, got {other:?}"),
        }
    }

    #[test]
    fn estimate_cost_requires_dry_run() {
        let err =
            Cli::try_parse_from(["mcptest", "pipe", "flow.yaml", "--estimate-cost"]).unwrap_err();
        assert!(err.to_string().contains("dry-run") || err.to_string().contains("dry_run"));
    }

    #[test]
    fn parses_on_budget_exceeded_continue() {
        let cli = Cli::try_parse_from([
            "mcptest",
            "pipe",
            "flow.yaml",
            "--on-budget-exceeded",
            "continue",
        ])
        .expect("parse");
        match cli.command {
            Command::Pipe(args) => {
                assert_eq!(args.on_budget_exceeded, OnBudgetExceededCli::Continue);
            }
            other => panic!("expected Pipe, got {other:?}"),
        }
    }
}
