//! Provider pricing table for cost projection.
//!
//! Cost projection (`mcptest pipe --dry-run --estimate-cost`) and the per-step
//! breakdown need to turn a model name plus a token estimate into dollars.
//! That mapping changes whenever a provider updates its prices, so it lives in
//! a data file (`pricing.yaml`) rather than in code: a price change is a YAML
//! edit, not a release. A bundled copy compiled into the binary is the
//! default; `--pricing-table <path>` swaps in a local override.

use std::collections::HashSet;
use std::sync::{Mutex, OnceLock};

use serde::Deserialize;
use thiserror::Error;

/// Bundled pricing table compiled into the binary so a release always has a
/// usable default even when no local `pricing.yaml` is present. Vendored
/// in-crate (kept in sync with the canonical workspace-root `pricing.yaml` by
/// `scripts/check-vendored-assets.sh`) so `cargo publish` packages it.
const BUNDLED_PRICING: &str = include_str!("../../pricing.yaml");

/// Errors raised while loading a pricing table.
#[derive(Debug, Error)]
pub enum PricingError {
    /// The override file could not be read.
    #[error("could not read pricing table `{path}`: {source}")]
    Read {
        /// The path we tried to read.
        path: String,
        /// The underlying IO error.
        source: std::io::Error,
    },
    /// The pricing YAML did not parse.
    #[error("pricing table `{path}` is not valid YAML: {source}")]
    Parse {
        /// The offending file.
        path: String,
        /// The YAML parse error.
        source: serde_yaml::Error,
    },
}

/// One model's token rates, USD per million tokens.
#[derive(Debug, Clone, Deserialize)]
pub struct ModelRate {
    /// Case-insensitive substring matched against the model name.
    #[serde(rename = "match")]
    pub match_substring: String,
    /// Input (prompt) rate, USD per million tokens.
    pub input_per_million: f64,
    /// Output (completion) rate, USD per million tokens.
    pub output_per_million: f64,
}

/// Default rates applied when no model row matches.
#[derive(Debug, Clone, Deserialize)]
pub struct DefaultRate {
    /// Input rate, USD per million tokens.
    pub input_per_million: f64,
    /// Output rate, USD per million tokens.
    pub output_per_million: f64,
}

/// The full pricing table: per-model rates plus a non-LLM tool-call cost.
#[derive(Debug, Clone, Deserialize)]
pub struct PricingTable {
    /// Flat USD cost attributed to a tool call that does not touch an LLM.
    pub tool_call_usd: f64,
    /// Per-model token rates, consulted longest-match-first.
    pub models: Vec<ModelRate>,
    /// Fallback rates when no model row matches.
    pub default: DefaultRate,
}

impl PricingTable {
    /// Load the bundled default table.
    ///
    /// Infallible in practice (the file ships with the binary) but returns
    /// `Result` so a malformed bundled file surfaces as an error rather than
    /// a panic during a release smoke test.
    pub fn bundled() -> Result<Self, PricingError> {
        serde_yaml::from_str(BUNDLED_PRICING).map_err(|source| PricingError::Parse {
            path: "<bundled pricing.yaml>".to_string(),
            source,
        })
    }

    /// Load a user-supplied override table from disk.
    pub fn load(path: &std::path::Path) -> Result<Self, PricingError> {
        let display = path.display().to_string();
        let raw = std::fs::read_to_string(path).map_err(|source| PricingError::Read {
            path: display.clone(),
            source,
        })?;
        serde_yaml::from_str(&raw).map_err(|source| PricingError::Parse {
            path: display,
            source,
        })
    }

    /// Per-call cost for a plain tool call (no LLM involved).
    pub fn tool_call_cost(&self) -> f64 {
        self.tool_call_usd
    }

    /// Estimate the dollar cost of `input_tokens` and `output_tokens` for a
    /// given model name.
    ///
    /// The model rows are consulted longest-match-first so a specific name
    /// (`claude-3-opus`) wins over a generic substring (`opus`) when both
    /// appear in the table. An unmatched model falls back to the default
    /// rates, which the bundled file sets conservatively high, and logs a
    /// once-per-model warning naming the model and the rates used. The bare
    /// `f64` return cannot carry a "this is a default-rate guess" flag
    /// without breaking existing callers, so callers that want to label the
    /// figure ask [`Self::is_default_rate`] instead.
    pub fn estimate(&self, model: &str, input_tokens: u64, output_tokens: u64) -> f64 {
        let (input_rate, output_rate) = match self.rate_for(model) {
            Some(row) => (row.input_per_million, row.output_per_million),
            None => {
                warn_default_rate_once(
                    model,
                    self.default.input_per_million,
                    self.default.output_per_million,
                );
                (
                    self.default.input_per_million,
                    self.default.output_per_million,
                )
            }
        };
        (input_tokens as f64) * input_rate / 1_000_000.0
            + (output_tokens as f64) * output_rate / 1_000_000.0
    }

    /// True when [`Self::estimate`] would price `model` with the table's
    /// default rates because no row matches. Exposed so reporters can mark
    /// such figures as default-rate guesses; `estimate` itself returns a
    /// bare `f64` and cannot carry the flag without breaking its callers.
    pub fn is_default_rate(&self, model: &str) -> bool {
        self.rate_for(model).is_none()
    }

    /// Longest-substring-match row lookup shared by `estimate` and
    /// `is_default_rate`, so both agree on what "unknown model" means.
    fn rate_for(&self, model: &str) -> Option<&ModelRate> {
        let lower = model.to_ascii_lowercase();
        let mut best: Option<&ModelRate> = None;
        for row in &self.models {
            if lower.contains(&row.match_substring.to_ascii_lowercase()) {
                let longer = best
                    .map(|b| row.match_substring.len() > b.match_substring.len())
                    .unwrap_or(true);
                if longer {
                    best = Some(row);
                }
            }
        }
        best
    }
}

/// Warn that `model` is being priced at default rates, at most once per
/// model per process so a pipe estimating hundreds of steps does not
/// repeat itself. Returns whether this call logged, so the once-only
/// behavior is testable without installing a tracing subscriber.
fn warn_default_rate_once(model: &str, input_per_million: f64, output_per_million: f64) -> bool {
    static WARNED: OnceLock<Mutex<HashSet<String>>> = OnceLock::new();
    let registry = WARNED.get_or_init(|| Mutex::new(HashSet::new()));
    let Ok(mut warned) = registry.lock() else {
        // A poisoned registry only loses log deduplication; pricing must
        // keep working, so skip the warning rather than fail or panic.
        return false;
    };
    if !warned.insert(model.to_string()) {
        return false;
    }
    tracing::warn!(
        target: "mcptest::pipe::pricing",
        model,
        input_per_million,
        output_per_million,
        "model not found in the pricing table; estimate uses default rates and may be inaccurate"
    );
    true
}

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

    #[test]
    fn bundled_table_loads() {
        let table = PricingTable::bundled().unwrap();
        assert!(table.tool_call_cost() >= 0.0);
        assert!(!table.models.is_empty());
    }

    #[test]
    fn longest_match_wins_for_opus() {
        let table = PricingTable::bundled().unwrap();
        // claude-3-opus and opus both match; the more specific row should win,
        // but both carry the same rate here so we just confirm a sane number.
        let cost = table.estimate("claude-3-opus-20240229", 1_000_000, 0);
        assert!(
            (cost - 15.0).abs() < 1e-6,
            "expected $15 for 1M opus input, got {cost}"
        );
    }

    #[test]
    fn sonnet_rates_apply_to_sonnet_models() {
        let table = PricingTable::bundled().unwrap();
        let cost = table.estimate("claude-3-5-sonnet-20241022", 1_000_000, 1_000_000);
        // 3 + 15 = 18 dollars for 1M in + 1M out.
        assert!((cost - 18.0).abs() < 1e-6, "got {cost}");
    }

    #[test]
    fn unknown_model_falls_back_to_default() {
        let table = PricingTable::bundled().unwrap();
        let cost = table.estimate("some-unlisted-model", 1_000_000, 0);
        assert!(
            (cost - 3.0).abs() < 1e-6,
            "expected default $3 input rate, got {cost}"
        );
    }

    #[test]
    fn unknown_model_warns_once_per_model() {
        // The registry is process-wide, so use names no other test prices.
        assert!(warn_default_rate_once("warn-once-test-model", 3.0, 15.0));
        assert!(
            !warn_default_rate_once("warn-once-test-model", 3.0, 15.0),
            "second estimate for the same model must not warn again"
        );
        assert!(
            warn_default_rate_once("warn-once-other-model", 3.0, 15.0),
            "a different model gets its own warning"
        );
    }

    #[test]
    fn is_default_rate_distinguishes_listed_from_unlisted_models() {
        let table = PricingTable::bundled().unwrap();
        assert!(!table.is_default_rate("claude-3-5-sonnet-20241022"));
        assert!(table.is_default_rate("some-unlisted-model"));
    }

    #[test]
    fn zero_tokens_cost_nothing() {
        let table = PricingTable::bundled().unwrap();
        assert_eq!(table.estimate("sonnet", 0, 0), 0.0);
    }
}
