//! Handler for `mcptest pipe` and its budget controls.
//!
//! Wires the clap args to the pipeline engine in [`crate::pipe`]: load and
//! validate the YAML, branch on `--dry-run` (plan + optional cost projection,
//! no connection), otherwise connect, run the steps under the budget, and
//! render output to the right stream. The result (and any budget breach)
//! decides the exit code.

use std::collections::HashMap;

use anyhow::{Context, Result};

use crate::cli::args::{OnBudgetExceededCli, PipeArgs, PipeFormat};
use crate::pipe::budget::{BudgetLimits, BudgetTracker, OnBudgetExceeded};
use crate::pipe::conn::Connection;
use crate::pipe::model::Pipeline;
use crate::pipe::plan::project;
use crate::pipe::pricing::PricingTable;
use crate::pipe::{executor, report};

/// Run a pipeline file end to end.
pub(crate) fn pipe_command(args: PipeArgs) -> Result<i32> {
    let pipeline = Pipeline::load(&args.pipeline)
        .with_context(|| format!("loading pipeline `{}`", args.pipeline.display()))?;
    let pricing = load_pricing(&args)?;

    if args.dry_run {
        // No tool calls in dry-run; v0 OSS pipelines are plain tool calls so
        // every step is costed at the per-call rate (model_for returns None).
        let projection = project(&pipeline, &pricing, |_| None);
        let display = args.pipeline.display().to_string();
        print!(
            "{}",
            report::render_dry_run(&display, &projection, args.estimate_cost)
        );
        return Ok(0);
    }

    let url = args
        .url
        .clone()
        .ok_or_else(|| anyhow::anyhow!("`pipe` requires --url to run (use --dry-run otherwise)"))?;
    let bearer_token = args
        .bearer_token_env
        .as_ref()
        .and_then(|name| std::env::var(name).ok());
    let vars: HashMap<String, String> = args.vars.iter().cloned().collect();
    let budget = BudgetTracker::new(build_limits(&args));

    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .context("failed to build tokio runtime")?;
    runtime.block_on(async move {
        let conn = Connection::connect(&url, bearer_token).await?;
        let result = executor::run(&pipeline, &conn, &vars, budget, &pricing).await;
        conn.close().await;
        emit_result(&result, args.format);
        Ok(exit_code(&result))
    })
}

fn load_pricing(args: &PipeArgs) -> Result<PricingTable> {
    match &args.pricing_table {
        Some(path) => PricingTable::load(path)
            .with_context(|| format!("loading pricing table `{}`", path.display())),
        None => PricingTable::bundled().context("loading bundled pricing table"),
    }
}

fn build_limits(args: &PipeArgs) -> BudgetLimits {
    BudgetLimits {
        max_cost: args.max_cost,
        max_tokens: args.max_tokens,
        max_cost_per_call: args.max_cost_per_call,
        max_duration: args.max_duration,
        on_exceeded: match args.on_budget_exceeded {
            OnBudgetExceededCli::Stop => OnBudgetExceeded::Stop,
            OnBudgetExceededCli::Continue => OnBudgetExceeded::Continue,
            OnBudgetExceededCli::Warn => OnBudgetExceeded::Warn,
        },
    }
}

/// Print the result: the chosen format goes to stdout (so it pipes cleanly);
/// the per-step cost breakdown and any breach message go to stderr.
fn emit_result(result: &executor::RunResult, format: PipeFormat) {
    match format {
        PipeFormat::Pretty => println!("{}", report::render_last_output_pretty(result)),
        PipeFormat::Json => {
            let trace = report::render_trace_json(result);
            println!(
                "{}",
                serde_json::to_string_pretty(&trace).unwrap_or_else(|_| trace.to_string())
            );
        }
    }
    if result.budget.steps().iter().any(|s| s.cost_usd > 0.0)
        || result.budget.cost_cap().is_some()
        || result.budget.token_cap().is_some()
    {
        eprint!("{}", report::render_cost_breakdown(&result.budget));
    }
    report_failures(result);
}

fn report_failures(result: &executor::RunResult) {
    for step in &result.steps {
        if let crate::pipe::executor::StepStatus::FailedContinued(msg) = &step.status {
            eprintln!("step `{}` failed: {msg}", step.id);
        }
    }
    if let Some((step_id, kind)) = &result.breach {
        eprintln!("budget exceeded ({} cap) at step `{step_id}`", kind.label());
    }
}

fn exit_code(result: &executor::RunResult) -> i32 {
    if result.is_failure() {
        1
    } else {
        0
    }
}
