//! Dry-run planning and cost projection ( `--dry-run`,
//! `--estimate-cost`).
//!
//! A dry run answers two questions without making a single paid call: what
//! would this pipeline do, and roughly what would it cost? The plan lists each
//! step in order; the cost projection attributes a per-call cost to plain tool
//! calls and a token-based estimate to LLM-touching steps, then sums them and
//! adds a worst-case figure that accounts for retries and a token-estimate
//! safety margin. This is the CI sanity check the budget ticket calls for.

use super::model::{OnError, Pipeline, Step};
use super::pricing::PricingTable;

/// A projected cost line for one step.
#[derive(Debug, Clone)]
pub struct StepProjection {
    /// Step id.
    pub id: String,
    /// Estimated USD for the step.
    pub cost_usd: f64,
    /// One-line description of how the estimate was derived.
    pub basis: String,
}

/// The full projection for a pipeline.
#[derive(Debug, Clone)]
pub struct Projection {
    /// Per-step estimates in order.
    pub steps: Vec<StepProjection>,
    /// Sum of the per-step estimates.
    pub total_usd: f64,
    /// Worst-case total: retries applied and a 1.5x token-estimate margin.
    pub worst_case_usd: f64,
}

/// Heuristic LLM token estimate used when a step looks like it drives a
/// model. A pure tool-call pipeline never trips this; it exists so a
/// projection of an LLM-bearing step is not silently zero.
const LLM_INPUT_TOKEN_ESTIMATE: u64 = 8_200;
/// Companion output estimate for [`LLM_INPUT_TOKEN_ESTIMATE`].
const LLM_OUTPUT_TOKEN_ESTIMATE: u64 = 2_000;
/// Worst-case multiplier on the token estimate.
const WORST_CASE_TOKEN_MULTIPLIER: f64 = 1.5;
/// Worst-case retry count assumed per step.
const WORST_CASE_RETRIES: u32 = 3;

/// Build a cost projection for a pipeline.
///
/// `model_for` maps a step to the model name it would drive, when the step is
/// LLM-bearing. Returning `None` means the step is a plain tool call and is
/// costed at the table's per-call rate. v0 has no LLM steps in OSS pipelines,
/// so the executor passes a closure that always returns `None`; the parameter
/// exists so the agent-test integration can supply real model names.
pub fn project<F>(pipeline: &Pipeline, pricing: &PricingTable, model_for: F) -> Projection
where
    F: Fn(&Step) -> Option<String>,
{
    let mut steps = Vec::with_capacity(pipeline.steps.len());
    let mut total = 0.0;
    let mut worst = 0.0;
    for step in &pipeline.steps {
        let (cost, basis) = match model_for(step) {
            Some(model) => {
                let c =
                    pricing.estimate(&model, LLM_INPUT_TOKEN_ESTIMATE, LLM_OUTPUT_TOKEN_ESTIMATE);
                (
                    c,
                    format!(
                        "LLM call: {model}, est. {LLM_INPUT_TOKEN_ESTIMATE} input + {LLM_OUTPUT_TOKEN_ESTIMATE} output tokens"
                    ),
                )
            }
            None => (pricing.tool_call_cost(), "tool call, no LLM".to_string()),
        };
        total += cost;
        worst += worst_case_step(step, cost);
        steps.push(StepProjection {
            id: step.id.clone(),
            cost_usd: cost,
            basis,
        });
    }
    Projection {
        steps,
        total_usd: total,
        worst_case_usd: worst,
    }
}

fn worst_case_step(step: &Step, base_cost: f64) -> f64 {
    let retries = match step.on_error {
        OnError::Retry(n) => n.max(WORST_CASE_RETRIES),
        _ => WORST_CASE_RETRIES,
    };
    base_cost * WORST_CASE_TOKEN_MULTIPLIER * f64::from(retries + 1)
}

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

    fn pipeline_with(n: usize) -> Pipeline {
        let yaml = {
            let mut s = String::from("pipe:\n");
            for i in 0..n {
                s.push_str(&format!("  - id: s{i}\n    tool: t\n"));
            }
            s
        };
        Pipeline::from_yaml_str(&yaml, "<test>").unwrap()
    }

    #[test]
    fn projects_tool_calls_at_per_call_rate() {
        let pricing = PricingTable::bundled().unwrap();
        let p = pipeline_with(3);
        let projection = project(&p, &pricing, |_| None);
        assert_eq!(projection.steps.len(), 3);
        let expected = pricing.tool_call_cost() * 3.0;
        assert!((projection.total_usd - expected).abs() < 1e-9);
        assert!(projection.worst_case_usd > projection.total_usd);
    }

    #[test]
    fn projects_llm_step_from_pricing_table() {
        let pricing = PricingTable::bundled().unwrap();
        let p = pipeline_with(1);
        let projection = project(&p, &pricing, |_| Some("claude-3-5-sonnet".to_string()));
        // 8200 input + 2000 output at sonnet rates is well above a tool call.
        assert!(projection.total_usd > pricing.tool_call_cost());
        assert!(projection.steps[0].basis.contains("LLM call"));
    }

    #[test]
    fn worst_case_exceeds_nominal() {
        let pricing = PricingTable::bundled().unwrap();
        let p = pipeline_with(2);
        let projection = project(&p, &pricing, |_| None);
        assert!(projection.worst_case_usd > projection.total_usd);
    }
}
