//! Cumulative budget controls for a pipeline run.
//!
//! Budget is a property of a whole workflow, not a single call. A pipeline can
//! issue N tool calls, each spending money or tokens, so the cap that matters
//! is the running total across every step. This module tracks that total and
//! decides, before each step, whether running it would breach the cap. It also
//! records the per-step cost/token/duration breakdown that every reporter
//! surfaces.
//!
//! It deliberately mirrors the shape of [`mcptest_core::eval::cost::CostBudget`]
//! (the single-run `--max-cost`) rather than reinventing it: same
//! "record then check would_exceed" rhythm, extended to tokens and wall-clock
//! and to a per-step ledger.

use std::time::Duration;

/// How the pipeline reacts when a budget would be (or has been) breached.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum OnBudgetExceeded {
    /// Stop before the breaching step and exit non-zero (default).
    #[default]
    Stop,
    /// Run every step, then exit non-zero if the total exceeded a cap.
    Continue,
    /// Run every step, warn on stderr, but exit zero.
    Warn,
}

/// The set of caps a run was launched with. `None` means that dimension is
/// unbounded.
#[derive(Debug, Clone, Copy, Default)]
pub struct BudgetLimits {
    /// Aggregate USD ceiling across all steps.
    pub max_cost: Option<f64>,
    /// Aggregate token ceiling (input + output) across all steps.
    pub max_tokens: Option<u64>,
    /// Per-step USD ceiling, layered on top of the aggregate cap.
    pub max_cost_per_call: Option<f64>,
    /// Wall-clock ceiling for the whole run.
    pub max_duration: Option<Duration>,
    /// What to do on breach.
    pub on_exceeded: OnBudgetExceeded,
}

impl BudgetLimits {
    /// Whether any cap is set. When false the budget machinery is inert and
    /// the run behaves exactly as it did before .
    pub fn any_set(&self) -> bool {
        self.max_cost.is_some()
            || self.max_tokens.is_some()
            || self.max_cost_per_call.is_some()
            || self.max_duration.is_some()
    }
}

/// The cost, token, and timing record for a single executed step.
#[derive(Debug, Clone)]
pub struct StepCost {
    /// Step id.
    pub step_id: String,
    /// USD attributed to this step.
    pub cost_usd: f64,
    /// Input tokens, when the step touched an LLM.
    pub input_tokens: u64,
    /// Output tokens, when the step touched an LLM.
    pub output_tokens: u64,
    /// Wall-clock duration of the step.
    pub duration: Duration,
    /// True when the step touched an LLM (so reporters can show token
    /// columns only where they are meaningful).
    pub is_llm: bool,
}

impl StepCost {
    /// Total tokens (input + output) for this step.
    pub fn total_tokens(&self) -> u64 {
        self.input_tokens + self.output_tokens
    }
}

/// Reason a [`BudgetTracker`] refused to run, or flagged after a run.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum BreachKind {
    /// The aggregate USD cap would be (or was) exceeded.
    Cost,
    /// The aggregate token cap would be (or was) exceeded.
    Tokens,
    /// A single step exceeded the per-call USD cap.
    PerCallCost,
    /// The wall-clock cap was exceeded.
    Duration,
}

impl BreachKind {
    /// Human-readable label used in the stderr breach message.
    pub fn label(&self) -> &'static str {
        match self {
            BreachKind::Cost => "cost",
            BreachKind::Tokens => "token",
            BreachKind::PerCallCost => "per-call cost",
            BreachKind::Duration => "duration",
        }
    }
}

/// Running ledger that gates each step and rolls up the final breakdown.
#[derive(Debug, Clone)]
pub struct BudgetTracker {
    limits: BudgetLimits,
    spent_usd: f64,
    spent_tokens: u64,
    elapsed: Duration,
    steps: Vec<StepCost>,
}

impl BudgetTracker {
    /// Build a tracker for the given caps.
    pub fn new(limits: BudgetLimits) -> Self {
        Self {
            limits,
            spent_usd: 0.0,
            spent_tokens: 0,
            elapsed: Duration::ZERO,
            steps: Vec::new(),
        }
    }

    /// Whether running a step that would add `next` cost and `next_tokens`
    /// tokens would breach a cap, given the time already elapsed.
    ///
    /// Returns the first dimension that would breach so the caller can
    /// report it precisely. A `1e-9` slop on the cost comparison absorbs
    /// floating-point noise so a projected total that lands exactly on the
    /// cap is allowed.
    pub fn would_exceed(
        &self,
        next_cost: f64,
        next_tokens: u64,
        next_per_call: f64,
    ) -> Option<BreachKind> {
        if let Some(cap) = self.limits.max_cost_per_call {
            if next_per_call > cap + 1e-9 {
                return Some(BreachKind::PerCallCost);
            }
        }
        if let Some(cap) = self.limits.max_cost {
            if self.spent_usd + next_cost > cap + 1e-9 {
                return Some(BreachKind::Cost);
            }
        }
        if let Some(cap) = self.limits.max_tokens {
            if self.spent_tokens + next_tokens > cap {
                return Some(BreachKind::Tokens);
            }
        }
        if let Some(cap) = self.limits.max_duration {
            if self.elapsed >= cap {
                return Some(BreachKind::Duration);
            }
        }
        None
    }

    /// Record a completed step's actuals into the ledger.
    pub fn record(&mut self, step: StepCost) {
        self.spent_usd += step.cost_usd;
        self.spent_tokens += step.total_tokens();
        self.elapsed += step.duration;
        self.steps.push(step);
    }

    /// After a `continue`/`warn` run, the kind of cap that ended up breached,
    /// if any. Used by the reporter and the exit-code decision.
    pub fn post_run_breach(&self) -> Option<BreachKind> {
        if let Some(cap) = self.limits.max_cost {
            if self.spent_usd > cap + 1e-9 {
                return Some(BreachKind::Cost);
            }
        }
        if let Some(cap) = self.limits.max_tokens {
            if self.spent_tokens > cap {
                return Some(BreachKind::Tokens);
            }
        }
        if let Some(cap) = self.limits.max_duration {
            if self.elapsed > cap {
                return Some(BreachKind::Duration);
            }
        }
        None
    }

    /// Configured behavior on breach.
    pub fn on_exceeded(&self) -> OnBudgetExceeded {
        self.limits.on_exceeded
    }

    /// Total dollars spent so far.
    pub fn spent_usd(&self) -> f64 {
        self.spent_usd
    }

    /// Total tokens spent so far.
    pub fn spent_tokens(&self) -> u64 {
        self.spent_tokens
    }

    /// Configured aggregate cost cap, if any.
    pub fn cost_cap(&self) -> Option<f64> {
        self.limits.max_cost
    }

    /// Configured aggregate token cap, if any.
    pub fn token_cap(&self) -> Option<u64> {
        self.limits.max_tokens
    }

    /// Per-step breakdown in execution order.
    pub fn steps(&self) -> &[StepCost] {
        &self.steps
    }

    /// Fraction (0.0..=1.0+) of the cost cap consumed, for the "17% used"
    /// reporter line. `None` when no cost cap is set.
    pub fn cost_fraction_used(&self) -> Option<f64> {
        self.limits.max_cost.map(|cap| {
            if cap <= 0.0 {
                1.0
            } else {
                self.spent_usd / cap
            }
        })
    }
}

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

    fn llm_step(id: &str, cost: f64, input: u64, output: u64) -> StepCost {
        StepCost {
            step_id: id.to_string(),
            cost_usd: cost,
            input_tokens: input,
            output_tokens: output,
            duration: Duration::from_millis(100),
            is_llm: true,
        }
    }

    #[test]
    fn unbounded_budget_never_blocks() {
        let t = BudgetTracker::new(BudgetLimits::default());
        assert_eq!(t.would_exceed(1_000.0, 1_000_000, 1_000.0), None);
        assert!(!t.limits.any_set());
    }

    #[test]
    fn cost_cap_blocks_the_step_that_would_overshoot() {
        let mut t = BudgetTracker::new(BudgetLimits {
            max_cost: Some(0.50),
            ..Default::default()
        });
        t.record(llm_step("a", 0.40, 100, 50));
        assert_eq!(t.would_exceed(0.05, 0, 0.05), None);
        assert_eq!(t.would_exceed(0.20, 0, 0.20), Some(BreachKind::Cost));
    }

    #[test]
    fn token_cap_blocks_on_token_total() {
        let mut t = BudgetTracker::new(BudgetLimits {
            max_tokens: Some(1_000),
            ..Default::default()
        });
        t.record(llm_step("a", 0.0, 600, 200));
        assert_eq!(t.would_exceed(0.0, 100, 0.0), None);
        assert_eq!(t.would_exceed(0.0, 300, 0.0), Some(BreachKind::Tokens));
    }

    #[test]
    fn per_call_cap_blocks_a_single_expensive_step() {
        let t = BudgetTracker::new(BudgetLimits {
            max_cost: Some(100.0),
            max_cost_per_call: Some(0.10),
            ..Default::default()
        });
        assert_eq!(t.would_exceed(0.20, 0, 0.20), Some(BreachKind::PerCallCost));
        assert_eq!(t.would_exceed(0.05, 0, 0.05), None);
    }

    #[test]
    fn post_run_breach_detects_overspend() {
        let mut t = BudgetTracker::new(BudgetLimits {
            max_cost: Some(0.10),
            on_exceeded: OnBudgetExceeded::Continue,
            ..Default::default()
        });
        t.record(llm_step("a", 0.30, 0, 0));
        assert_eq!(t.post_run_breach(), Some(BreachKind::Cost));
    }

    #[test]
    fn cost_fraction_tracks_usage() {
        let mut t = BudgetTracker::new(BudgetLimits {
            max_cost: Some(0.50),
            ..Default::default()
        });
        t.record(llm_step("a", 0.085, 0, 0));
        let frac = t.cost_fraction_used().unwrap();
        assert!((frac - 0.17).abs() < 0.01, "expected ~0.17, got {frac}");
    }

    #[test]
    fn steps_are_recorded_in_order() {
        let mut t = BudgetTracker::new(BudgetLimits::default());
        t.record(llm_step("first", 0.01, 0, 0));
        t.record(llm_step("second", 0.02, 0, 0));
        let ids: Vec<&str> = t.steps().iter().map(|s| s.step_id.as_str()).collect();
        assert_eq!(ids, vec!["first", "second"]);
    }

    #[test]
    fn exactly_on_cap_is_allowed() {
        let mut t = BudgetTracker::new(BudgetLimits {
            max_cost: Some(0.50),
            ..Default::default()
        });
        t.record(llm_step("a", 0.30, 0, 0));
        // 0.30 + 0.20 == 0.50 exactly; must not breach.
        assert_eq!(t.would_exceed(0.20, 0, 0.20), None);
    }
}
