swink-agent 0.13.2

use std::sync::OnceLock;

use chrono::NaiveDate;
use serde::Deserialize;

use crate::ModelSpec;
use crate::pricing::CostCalculator;
use crate::types::{
    AssistantMessage, Cost, ModelCapabilities, ThinkingLevel, ThinkingLevelSet, Usage,
};

/// Whether a provider's models run on a remote API or on local hardware.
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ProviderKind {
    Remote,
    Local,
}

/// How requests to a provider are authenticated.
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AuthMode {
    Bearer,
    ApiKeyHeader,
    AwsSigv4,
}

/// Provider API version selector used when building request URLs.
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ApiVersion {
    V1,
    V1beta,
}

/// A capability a preset's model supports, as declared in the catalog TOML.
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PresetCapability {
    Text,
    Tools,
    Thinking,
    ImagesIn,
    Streaming,
    StructuredOutput,
}

/// Release maturity of a preset's model at the provider.
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PresetStatus {
    Ga,
    Preview,
    /// The provider has retired (or announced retirement of) this model.
    ///
    /// The preset stays listed so catalog lookups and cost calculation keep
    /// working for historical data, and `replacement_model_id` points
    /// consumers at the successor model when one is known.
    ///
    /// TOML representation (existing string statuses are unaffected):
    ///
    /// ```toml
    /// [providers.presets.status.deprecated]
    /// replacement_model_id = "gpt-5.4"
    /// ```
    Deprecated {
        #[serde(default)]
        replacement_model_id: Option<String>,
    },
}

impl PresetStatus {
    /// Returns `true` for [`PresetStatus::Deprecated`], regardless of whether
    /// a replacement model is recorded.
    #[must_use]
    pub const fn is_deprecated(&self) -> bool {
        matches!(self, Self::Deprecated { .. })
    }
}

/// A single named model preset within a [`ProviderCatalog`], as loaded from the catalog TOML.
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Deserialize)]
pub struct PresetCatalog {
    pub id: String,
    pub display_name: String,
    pub group: Option<String>,
    pub model_id: String,
    pub api_version: Option<ApiVersion>,
    #[serde(default)]
    pub capabilities: Vec<PresetCapability>,
    pub status: Option<PresetStatus>,
    pub context_window_tokens: Option<u64>,
    pub max_output_tokens: Option<u64>,
    #[serde(default)]
    pub include_by_default: bool,
    pub repo_id: Option<String>,
    pub filename: Option<String>,
    #[serde(default)]
    pub cost_per_million_input: Option<f64>,
    #[serde(default)]
    pub cost_per_million_output: Option<f64>,
    #[serde(default)]
    pub cost_per_million_cache_read: Option<f64>,
    #[serde(default)]
    pub cost_per_million_cache_write: Option<f64>,
    /// Reasoning levels this model accepts. `None` = unannotated (falls back
    /// to `capabilities.contains(&PresetCapability::Thinking)`); `Some(&[])`
    /// = no reasoning-level control.
    pub reasoning_levels: Option<ThinkingLevelSet>,
}

impl PresetCatalog {
    /// Create a preset with the required identifying fields; everything else
    /// starts unset and can be filled in with the `with_*` builders.
    #[must_use]
    pub fn new(
        id: impl Into<String>,
        display_name: impl Into<String>,
        model_id: impl Into<String>,
    ) -> Self {
        Self {
            id: id.into(),
            display_name: display_name.into(),
            group: None,
            model_id: model_id.into(),
            api_version: None,
            capabilities: Vec::new(),
            status: None,
            context_window_tokens: None,
            max_output_tokens: None,
            include_by_default: false,
            repo_id: None,
            filename: None,
            cost_per_million_input: None,
            cost_per_million_output: None,
            cost_per_million_cache_read: None,
            cost_per_million_cache_write: None,
            reasoning_levels: None,
        }
    }

    /// Set the display group this preset is listed under.
    #[must_use]
    pub fn with_group(mut self, group: impl Into<String>) -> Self {
        self.group = Some(group.into());
        self
    }

    /// Set the provider API version used when building request URLs.
    #[must_use]
    pub fn with_api_version(mut self, api_version: ApiVersion) -> Self {
        self.api_version = Some(api_version);
        self
    }

    /// Set the declared capabilities.
    #[must_use]
    pub fn with_capabilities(mut self, capabilities: Vec<PresetCapability>) -> Self {
        self.capabilities = capabilities;
        self
    }

    /// Set the reasoning levels this model accepts.
    #[must_use]
    pub const fn with_reasoning_levels(mut self, levels: ThinkingLevelSet) -> Self {
        self.reasoning_levels = Some(levels);
        self
    }

    /// Set the release maturity status.
    #[must_use]
    pub fn with_status(mut self, status: PresetStatus) -> Self {
        self.status = Some(status);
        self
    }

    /// Set the model's context window size, in tokens.
    #[must_use]
    pub const fn with_context_window_tokens(mut self, tokens: u64) -> Self {
        self.context_window_tokens = Some(tokens);
        self
    }

    /// Set the model's maximum output tokens.
    #[must_use]
    pub const fn with_max_output_tokens(mut self, tokens: u64) -> Self {
        self.max_output_tokens = Some(tokens);
        self
    }

    /// Set whether this preset is included by default.
    #[must_use]
    pub const fn with_include_by_default(mut self, include: bool) -> Self {
        self.include_by_default = include;
        self
    }

    /// Set the local model repository identifier.
    #[must_use]
    pub fn with_repo_id(mut self, repo_id: impl Into<String>) -> Self {
        self.repo_id = Some(repo_id.into());
        self
    }

    /// Set the local model file name.
    #[must_use]
    pub fn with_filename(mut self, filename: impl Into<String>) -> Self {
        self.filename = Some(filename.into());
        self
    }

    /// Set the USD-per-million-input-token rate.
    #[must_use]
    pub const fn with_cost_per_million_input(mut self, cost: f64) -> Self {
        self.cost_per_million_input = Some(cost);
        self
    }

    /// Set the USD-per-million-output-token rate.
    #[must_use]
    pub const fn with_cost_per_million_output(mut self, cost: f64) -> Self {
        self.cost_per_million_output = Some(cost);
        self
    }

    /// Set the USD-per-million-cache-read-token rate.
    #[must_use]
    pub const fn with_cost_per_million_cache_read(mut self, cost: f64) -> Self {
        self.cost_per_million_cache_read = Some(cost);
        self
    }

    /// Set the USD-per-million-cache-write-token rate.
    #[must_use]
    pub const fn with_cost_per_million_cache_write(mut self, cost: f64) -> Self {
        self.cost_per_million_cache_write = Some(cost);
        self
    }
}

/// A provider entry in the model catalog, holding its auth/connection settings and presets.
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Deserialize)]
pub struct ProviderCatalog {
    pub key: String,
    pub display_name: String,
    pub kind: ProviderKind,
    pub auth_mode: Option<AuthMode>,
    pub credential_env_var: Option<String>,
    pub base_url_env_var: Option<String>,
    pub default_base_url: Option<String>,
    #[serde(default)]
    pub requires_base_url: bool,
    pub region_env_var: Option<String>,
    #[serde(default)]
    pub presets: Vec<PresetCatalog>,
}

impl ProviderCatalog {
    /// Create a provider entry with the required identifying fields; everything
    /// else starts unset and can be filled in with the `with_*` builders.
    #[must_use]
    pub fn new(
        key: impl Into<String>,
        display_name: impl Into<String>,
        kind: ProviderKind,
    ) -> Self {
        Self {
            key: key.into(),
            display_name: display_name.into(),
            kind,
            auth_mode: None,
            credential_env_var: None,
            base_url_env_var: None,
            default_base_url: None,
            requires_base_url: false,
            region_env_var: None,
            presets: Vec::new(),
        }
    }

    /// Set the authentication mode.
    #[must_use]
    pub fn with_auth_mode(mut self, auth_mode: AuthMode) -> Self {
        self.auth_mode = Some(auth_mode);
        self
    }

    /// Set the environment variable that holds the credential.
    #[must_use]
    pub fn with_credential_env_var(mut self, var: impl Into<String>) -> Self {
        self.credential_env_var = Some(var.into());
        self
    }

    /// Set the environment variable that holds the base URL override.
    #[must_use]
    pub fn with_base_url_env_var(mut self, var: impl Into<String>) -> Self {
        self.base_url_env_var = Some(var.into());
        self
    }

    /// Set the default base URL.
    #[must_use]
    pub fn with_default_base_url(mut self, url: impl Into<String>) -> Self {
        self.default_base_url = Some(url.into());
        self
    }

    /// Set whether a base URL is required to use this provider.
    #[must_use]
    pub const fn with_requires_base_url(mut self, requires: bool) -> Self {
        self.requires_base_url = requires;
        self
    }

    /// Set the environment variable that holds the region (e.g. for AWS).
    #[must_use]
    pub fn with_region_env_var(mut self, var: impl Into<String>) -> Self {
        self.region_env_var = Some(var.into());
        self
    }

    /// Set the provider's presets.
    #[must_use]
    pub fn with_presets(mut self, presets: Vec<PresetCatalog>) -> Self {
        self.presets = presets;
        self
    }

    #[must_use]
    pub fn preset(&self, preset_id: &str) -> Option<&PresetCatalog> {
        self.presets.iter().find(|preset| preset.id == preset_id)
    }
}

/// The full model catalog: a list of providers, each with its own presets.
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Deserialize)]
pub struct ModelCatalog {
    /// Date (`YYYY-MM-DD`) the compiled-in pricing table was last verified
    /// against provider published prices. Used by the pricing-staleness
    /// warning at agent construction; `None` disables the check.
    #[serde(default)]
    pub pricing_as_of: Option<String>,
    #[serde(default)]
    pub providers: Vec<ProviderCatalog>,
}

impl ModelCatalog {
    /// Create an empty catalog with no `pricing_as_of` date and no providers.
    #[must_use]
    pub fn new() -> Self {
        Self {
            pricing_as_of: None,
            providers: Vec::new(),
        }
    }

    /// Set the `pricing_as_of` date (`YYYY-MM-DD`).
    #[must_use]
    pub fn with_pricing_as_of(mut self, pricing_as_of: impl Into<String>) -> Self {
        self.pricing_as_of = Some(pricing_as_of.into());
        self
    }

    /// Set the catalog's providers.
    #[must_use]
    pub fn with_providers(mut self, providers: Vec<ProviderCatalog>) -> Self {
        self.providers = providers;
        self
    }

    #[must_use]
    pub fn provider(&self, provider_key: &str) -> Option<&ProviderCatalog> {
        self.providers
            .iter()
            .find(|provider| provider.key == provider_key)
    }

    /// Search across all providers for a preset matching the given `model_id`.
    #[must_use]
    pub fn find_preset_by_model_id(&self, model_id: &str) -> Option<CatalogPreset> {
        for provider in &self.providers {
            for preset in &provider.presets {
                if preset.model_id == model_id {
                    return self.preset(&provider.key, &preset.id);
                }
            }
        }
        None
    }

    /// Look up a preset by `model_id` under one specific provider.
    ///
    /// The same `model_id` can be listed under several providers at
    /// different prices (`gpt-5.6-luna` is metered under `openai` and free
    /// under the subscription-backed `codex`), so a provider-blind lookup
    /// returns whichever block comes first. Use this when the provider is
    /// known. Returns `None` if the provider has no such row — callers that
    /// want a fallback chain with [`find_preset_by_model_id`] themselves.
    ///
    /// [`find_preset_by_model_id`]: Self::find_preset_by_model_id
    #[must_use]
    pub fn find_preset(&self, provider_key: &str, model_id: &str) -> Option<CatalogPreset> {
        let provider = self.provider(provider_key)?;
        let preset = provider.presets.iter().find(|p| p.model_id == model_id)?;
        self.preset(&provider.key, &preset.id)
    }

    #[must_use]
    pub fn preset(&self, provider_key: &str, preset_id: &str) -> Option<CatalogPreset> {
        let provider = self.provider(provider_key)?;
        let preset = provider.preset(preset_id)?;
        Some(CatalogPreset {
            provider_key: provider.key.clone(),
            provider_display_name: provider.display_name.clone(),
            provider_kind: provider.kind.clone(),
            preset_id: preset.id.clone(),
            display_name: preset.display_name.clone(),
            group: preset.group.clone(),
            model_id: preset.model_id.clone(),
            api_version: preset.api_version.clone(),
            capabilities: preset.capabilities.clone(),
            status: preset.status.clone(),
            context_window_tokens: preset.context_window_tokens,
            max_output_tokens: preset.max_output_tokens,
            auth_mode: provider.auth_mode.clone(),
            credential_env_var: provider.credential_env_var.clone(),
            base_url_env_var: provider.base_url_env_var.clone(),
            default_base_url: provider.default_base_url.clone(),
            requires_base_url: provider.requires_base_url,
            region_env_var: provider.region_env_var.clone(),
            include_by_default: preset.include_by_default,
            repo_id: preset.repo_id.clone(),
            filename: preset.filename.clone(),
            cost_per_million_input: preset.cost_per_million_input,
            cost_per_million_output: preset.cost_per_million_output,
            cost_per_million_cache_read: preset.cost_per_million_cache_read,
            cost_per_million_cache_write: preset.cost_per_million_cache_write,
            reasoning_levels: preset.reasoning_levels,
        })
    }
}

impl Default for ModelCatalog {
    fn default() -> Self {
        Self::new()
    }
}

/// A preset flattened together with its parent provider's fields, for standalone use
/// once resolved via [`ModelCatalog::preset`].
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq)]
pub struct CatalogPreset {
    pub provider_key: String,
    pub provider_display_name: String,
    pub provider_kind: ProviderKind,
    pub preset_id: String,
    pub display_name: String,
    pub group: Option<String>,
    pub model_id: String,
    pub api_version: Option<ApiVersion>,
    pub capabilities: Vec<PresetCapability>,
    pub status: Option<PresetStatus>,
    pub context_window_tokens: Option<u64>,
    pub max_output_tokens: Option<u64>,
    pub auth_mode: Option<AuthMode>,
    pub credential_env_var: Option<String>,
    pub base_url_env_var: Option<String>,
    pub default_base_url: Option<String>,
    pub requires_base_url: bool,
    pub region_env_var: Option<String>,
    pub include_by_default: bool,
    pub repo_id: Option<String>,
    pub filename: Option<String>,
    pub cost_per_million_input: Option<f64>,
    pub cost_per_million_output: Option<f64>,
    pub cost_per_million_cache_read: Option<f64>,
    pub cost_per_million_cache_write: Option<f64>,
    /// Reasoning levels this model accepts. `None` = unannotated (falls back
    /// to `capabilities.contains(&PresetCapability::Thinking)`); `Some(&[])`
    /// = no reasoning-level control.
    pub reasoning_levels: Option<ThinkingLevelSet>,
}

impl CatalogPreset {
    /// Create a flattened preset with the required identifying fields;
    /// everything else starts unset and can be filled in with the `with_*`
    /// builders.
    #[must_use]
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        provider_key: impl Into<String>,
        provider_display_name: impl Into<String>,
        provider_kind: ProviderKind,
        preset_id: impl Into<String>,
        display_name: impl Into<String>,
        model_id: impl Into<String>,
    ) -> Self {
        Self {
            provider_key: provider_key.into(),
            provider_display_name: provider_display_name.into(),
            provider_kind,
            preset_id: preset_id.into(),
            display_name: display_name.into(),
            group: None,
            model_id: model_id.into(),
            api_version: None,
            capabilities: Vec::new(),
            status: None,
            context_window_tokens: None,
            max_output_tokens: None,
            auth_mode: None,
            credential_env_var: None,
            base_url_env_var: None,
            default_base_url: None,
            requires_base_url: false,
            region_env_var: None,
            include_by_default: false,
            repo_id: None,
            filename: None,
            cost_per_million_input: None,
            cost_per_million_output: None,
            cost_per_million_cache_read: None,
            cost_per_million_cache_write: None,
            reasoning_levels: None,
        }
    }

    /// Set the display group this preset is listed under.
    #[must_use]
    pub fn with_group(mut self, group: impl Into<String>) -> Self {
        self.group = Some(group.into());
        self
    }

    /// Set the provider API version used when building request URLs.
    #[must_use]
    pub fn with_api_version(mut self, api_version: ApiVersion) -> Self {
        self.api_version = Some(api_version);
        self
    }

    /// Set the declared capabilities.
    #[must_use]
    pub fn with_capabilities(mut self, capabilities: Vec<PresetCapability>) -> Self {
        self.capabilities = capabilities;
        self
    }

    /// Set the reasoning levels this model accepts.
    #[must_use]
    pub const fn with_reasoning_levels(mut self, levels: ThinkingLevelSet) -> Self {
        self.reasoning_levels = Some(levels);
        self
    }

    /// Set the release maturity status.
    #[must_use]
    pub fn with_status(mut self, status: PresetStatus) -> Self {
        self.status = Some(status);
        self
    }

    /// Set the model's context window size, in tokens.
    #[must_use]
    pub const fn with_context_window_tokens(mut self, tokens: u64) -> Self {
        self.context_window_tokens = Some(tokens);
        self
    }

    /// Set the model's maximum output tokens.
    #[must_use]
    pub const fn with_max_output_tokens(mut self, tokens: u64) -> Self {
        self.max_output_tokens = Some(tokens);
        self
    }

    /// Set the provider's authentication mode.
    #[must_use]
    pub fn with_auth_mode(mut self, auth_mode: AuthMode) -> Self {
        self.auth_mode = Some(auth_mode);
        self
    }

    /// Set the environment variable that holds the credential.
    #[must_use]
    pub fn with_credential_env_var(mut self, var: impl Into<String>) -> Self {
        self.credential_env_var = Some(var.into());
        self
    }

    /// Set the environment variable that holds the base URL override.
    #[must_use]
    pub fn with_base_url_env_var(mut self, var: impl Into<String>) -> Self {
        self.base_url_env_var = Some(var.into());
        self
    }

    /// Set the default base URL.
    #[must_use]
    pub fn with_default_base_url(mut self, url: impl Into<String>) -> Self {
        self.default_base_url = Some(url.into());
        self
    }

    /// Set whether a base URL is required to use this provider.
    #[must_use]
    pub const fn with_requires_base_url(mut self, requires: bool) -> Self {
        self.requires_base_url = requires;
        self
    }

    /// Set the environment variable that holds the region (e.g. for AWS).
    #[must_use]
    pub fn with_region_env_var(mut self, var: impl Into<String>) -> Self {
        self.region_env_var = Some(var.into());
        self
    }

    /// Set whether this preset is included by default.
    #[must_use]
    pub const fn with_include_by_default(mut self, include: bool) -> Self {
        self.include_by_default = include;
        self
    }

    /// Set the local model repository identifier.
    #[must_use]
    pub fn with_repo_id(mut self, repo_id: impl Into<String>) -> Self {
        self.repo_id = Some(repo_id.into());
        self
    }

    /// Set the local model file name.
    #[must_use]
    pub fn with_filename(mut self, filename: impl Into<String>) -> Self {
        self.filename = Some(filename.into());
        self
    }

    /// Set the USD-per-million-input-token rate.
    #[must_use]
    pub const fn with_cost_per_million_input(mut self, cost: f64) -> Self {
        self.cost_per_million_input = Some(cost);
        self
    }

    /// Set the USD-per-million-output-token rate.
    #[must_use]
    pub const fn with_cost_per_million_output(mut self, cost: f64) -> Self {
        self.cost_per_million_output = Some(cost);
        self
    }

    /// Set the USD-per-million-cache-read-token rate.
    #[must_use]
    pub const fn with_cost_per_million_cache_read(mut self, cost: f64) -> Self {
        self.cost_per_million_cache_read = Some(cost);
        self
    }

    /// Set the USD-per-million-cache-write-token rate.
    #[must_use]
    pub const fn with_cost_per_million_cache_write(mut self, cost: f64) -> Self {
        self.cost_per_million_cache_write = Some(cost);
        self
    }

    /// Build a [`ModelCapabilities`] from the catalog's capability list and
    /// token limits.
    #[must_use]
    pub fn model_capabilities(&self) -> ModelCapabilities {
        let has = |cap: &PresetCapability| self.capabilities.contains(cap);
        ModelCapabilities {
            supports_thinking: has(&PresetCapability::Thinking),
            supports_vision: has(&PresetCapability::ImagesIn),
            supports_tool_use: has(&PresetCapability::Tools),
            supports_streaming: has(&PresetCapability::Streaming),
            supports_structured_output: has(&PresetCapability::StructuredOutput),
            max_context_window: self.context_window_tokens,
            max_output_tokens: self.max_output_tokens,
            reasoning_levels: self.reasoning_levels,
        }
    }

    /// Create a [`ModelSpec`] pre-populated with capabilities from the catalog.
    ///
    /// Local thinking-capable models default to [`ThinkingLevel::Medium`] so
    /// thinking is active out of the box (local inference treats any non-`Off`
    /// level as a binary "on" toggle). Remote presets keep the opt-in
    /// [`ThinkingLevel::Off`] default because remote thinking consumes billable
    /// token budget. Callers can still disable thinking explicitly via
    /// [`ModelSpec::with_thinking_level`] with [`ThinkingLevel::Off`].
    #[must_use]
    pub fn model_spec(&self) -> ModelSpec {
        let capabilities = self.model_capabilities();
        let mut spec = ModelSpec::new(&self.provider_key, &self.model_id);
        if self.provider_kind == ProviderKind::Local && capabilities.supports_thinking {
            spec = spec.with_thinking_level(ThinkingLevel::Medium);
        }
        spec.with_capabilities(capabilities)
    }

    /// Returns `true` when the preset's status is [`PresetStatus::Deprecated`].
    #[must_use]
    pub fn is_deprecated(&self) -> bool {
        self.status
            .as_ref()
            .is_some_and(PresetStatus::is_deprecated)
    }

    /// The catalog-recorded replacement for a deprecated preset, if any.
    ///
    /// Returns `None` for non-deprecated presets and for deprecated presets
    /// without a known successor.
    #[must_use]
    pub fn replacement_model_id(&self) -> Option<&str> {
        match self.status.as_ref()? {
            PresetStatus::Deprecated {
                replacement_model_id,
            } => replacement_model_id.as_deref(),
            _ => None,
        }
    }
}

impl ModelCatalog {
    /// The parsed `pricing_as_of` date, or `None` if absent or malformed.
    #[must_use]
    pub fn pricing_as_of_date(&self) -> Option<NaiveDate> {
        NaiveDate::parse_from_str(self.pricing_as_of.as_deref()?, "%Y-%m-%d").ok()
    }

    /// Check whether the catalog's pricing data is stale as of `today`.
    ///
    /// Returns `Some(PricingStaleness)` when the pricing table is older than
    /// `threshold_days`, and `None` when it is fresh or when the catalog
    /// carries no (parseable) `pricing_as_of` date.
    #[must_use]
    pub fn pricing_staleness_at(
        &self,
        today: NaiveDate,
        threshold_days: u32,
    ) -> Option<PricingStaleness> {
        let as_of = self.pricing_as_of_date()?;
        let age_days = (today - as_of).num_days();
        (age_days > i64::from(threshold_days)).then_some(PricingStaleness {
            as_of,
            age_days,
            threshold_days,
        })
    }
}

/// Default staleness threshold (in days) for the compiled-in pricing table.
pub const DEFAULT_PRICING_STALENESS_DAYS: u32 = 180;

/// Environment variable that overrides [`DEFAULT_PRICING_STALENESS_DAYS`]
/// for the warning logged at agent construction. Value is a day count.
pub const PRICING_STALENESS_ENV_VAR: &str = "SWINK_PRICING_STALENESS_DAYS";

/// Details of a stale compiled-in pricing table.
///
/// Produced by [`pricing_staleness`] / [`ModelCatalog::pricing_staleness_at`]
/// when the catalog's `pricing_as_of` date is older than the threshold.
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PricingStaleness {
    /// Date the pricing table was last verified.
    pub as_of: NaiveDate,
    /// Age of the pricing table in days, relative to the evaluation date.
    pub age_days: i64,
    /// The threshold that was exceeded.
    pub threshold_days: u32,
}

impl PricingStaleness {
    /// Create a staleness record from its three fields.
    #[must_use]
    pub const fn new(as_of: NaiveDate, age_days: i64, threshold_days: u32) -> Self {
        Self {
            as_of,
            age_days,
            threshold_days,
        }
    }
}

/// Check the compiled-in catalog's pricing staleness against today's date.
///
/// Returns `Some` when the pricing table is older than `threshold_days`.
/// See [`DEFAULT_PRICING_STALENESS_DAYS`] for the default threshold used at
/// agent construction.
#[must_use]
pub fn pricing_staleness(threshold_days: u32) -> Option<PricingStaleness> {
    model_catalog().pricing_staleness_at(chrono::Utc::now().date_naive(), threshold_days)
}

/// Log a once-per-process warning when the compiled-in pricing table is
/// older than the configured threshold.
///
/// The threshold defaults to [`DEFAULT_PRICING_STALENESS_DAYS`] and can be
/// overridden via the [`PRICING_STALENESS_ENV_VAR`] environment variable.
/// Called at agent construction.
pub(crate) fn warn_if_pricing_stale() {
    static ONCE: std::sync::Once = std::sync::Once::new();
    ONCE.call_once(|| {
        let threshold_days = std::env::var(PRICING_STALENESS_ENV_VAR)
            .ok()
            .and_then(|value| value.trim().parse::<u32>().ok())
            .unwrap_or(DEFAULT_PRICING_STALENESS_DAYS);
        if let Some(staleness) = pricing_staleness(threshold_days) {
            tracing::warn!(
                pricing_as_of = %staleness.as_of,
                age_days = staleness.age_days,
                threshold_days = staleness.threshold_days,
                "compiled-in model pricing table may be stale; costs from \
                 calculate_cost() may not match current provider prices"
            );
        }
    });
}

#[must_use]
pub fn model_catalog() -> &'static ModelCatalog {
    static MODEL_CATALOG: OnceLock<ModelCatalog> = OnceLock::new();
    MODEL_CATALOG.get_or_init(|| {
        toml::from_str(include_str!("model_catalog.toml"))
            .expect("src/model_catalog.toml must be valid TOML")
    })
}

/// Compute monetary cost from token usage using catalog pricing data.
///
/// Looks up the model by `model_id` across all providers. Returns
/// `Cost::default()` if the model is not found or has no pricing data.
#[must_use]
pub fn calculate_cost(model_id: &str, usage: &Usage) -> Cost {
    let Some(preset) = model_catalog().find_preset_by_model_id(model_id) else {
        tracing::debug!(
            model_id,
            "model not found in catalog; cost reported as zero"
        );
        return Cost::default();
    };
    cost_from_preset(&preset, usage)
}

/// Like [`calculate_cost`], but prefers the rates listed under
/// `provider_key` and only falls back to the provider-blind lookup when
/// that provider has no row for `model_id`.
///
/// This is what keeps a subscription-backed provider (`codex`, priced at
/// zero) from being billed at the metered `openai` rates for the same slug.
#[must_use]
pub fn calculate_cost_for_provider(provider_key: &str, model_id: &str, usage: &Usage) -> Cost {
    match model_catalog().find_preset(provider_key, model_id) {
        Some(preset) => cost_from_preset(&preset, usage),
        None => calculate_cost(model_id, usage),
    }
}

fn cost_from_preset(preset: &CatalogPreset, usage: &Usage) -> Cost {
    #[allow(clippy::cast_precision_loss)] // token counts fit comfortably in f64
    let per_m = |tokens: u64, rate: Option<f64>| -> f64 {
        rate.map_or(0.0, |r| tokens as f64 * r / 1_000_000.0)
    };

    let input = per_m(usage.input, preset.cost_per_million_input);
    let output = per_m(usage.output, preset.cost_per_million_output);
    let cache_read = per_m(usage.cache_read, preset.cost_per_million_cache_read);
    let cache_write = per_m(usage.cache_write, preset.cost_per_million_cache_write);

    Cost {
        input,
        output,
        cache_read,
        cache_write,
        total: input + output + cache_read + cache_write,
        ..Cost::default()
    }
}

/// Fill in an assistant message's [`Cost`] from catalog pricing when the
/// adapter did not price the response itself.
///
/// Most built-in remote adapters emit `Cost::default()` on every assistant
/// message — they report token [`Usage`] but leave pricing to the caller. The
/// agent loop calls this helper on each assistant message before accumulating
/// cost, so that [`PolicyContext::accumulated_cost`](crate::PolicyContext) —
/// and therefore any cost ceiling built on it — sees real money.
///
/// Adapters that *do* supply their own cost (the proxy adapter, which passes
/// through provider-billed amounts) keep precedence: a non-zero [`Cost`] is
/// left untouched.
///
/// Returns `true` if the message was repriced, `false` if it was left as-is
/// (adapter already priced it, or the model has no catalog pricing).
///
/// # Example
/// ```rust
/// use swink_agent::{AssistantMessage, StopReason, Usage, price_assistant_message};
///
/// let mut message = AssistantMessage::new(vec![], "anthropic", "claude-sonnet-4-6")
///     .with_usage(Usage::default().with_input(1_000_000))
///     .with_stop_reason(StopReason::Stop)
///     .with_timestamp(0);
///
/// assert!(price_assistant_message(&mut message));
/// assert!((message.cost.total - 3.0).abs() < 1e-9);
/// ```
pub fn price_assistant_message(message: &mut AssistantMessage) -> bool {
    price_assistant_message_with(message, None)
}

/// Like [`price_assistant_message`], but consults an operator-declared
/// [`CostCalculator`] before falling back to the compiled model catalog.
///
/// This is what the agent loop actually calls, threading through the calculator
/// configured via
/// [`AgentOptions::with_cost_calculator`](crate::AgentOptions::with_cost_calculator).
/// It exists because the catalog only knows about models shipped with the
/// crate — local endpoints, private deployments, and negotiated per-tier rates
/// all price at zero without an override.
///
/// Precedence, highest first:
///
/// 1. The adapter's own non-zero [`Cost`] — never overwritten.
/// 2. `calculator`, when it returns a non-zero [`Cost`] for this model.
/// 3. The compiled model catalog.
///
/// Returns `true` if the message was repriced.
///
/// # Example
/// ```rust
/// use swink_agent::{
///     AssistantMessage, ModelRates, PricingTable, StopReason, Usage,
///     price_assistant_message_with,
/// };
///
/// // `claude-sonnet-4-6` is in the catalog at $3.00/M input, but the operator
/// // negotiated $1.00/M and says so.
/// let table = PricingTable::new().with_model(
///     "claude-sonnet-4-6",
///     ModelRates::default().with_input_per_million(1.0),
/// );
///
/// let mut message = AssistantMessage::new(vec![], "anthropic", "claude-sonnet-4-6")
///     .with_usage(Usage::default().with_input(1_000_000))
///     .with_stop_reason(StopReason::Stop)
///     .with_timestamp(0);
///
/// assert!(price_assistant_message_with(&mut message, Some(&table)));
/// assert!((message.cost.total - 1.0).abs() < 1e-9);
/// ```
pub fn price_assistant_message_with(
    message: &mut AssistantMessage,
    calculator: Option<&dyn CostCalculator>,
) -> bool {
    if !message.cost.is_zero() {
        return false;
    }
    let priced = calculator
        .and_then(|calculator| calculator.calculate(&message.model_id, &message.usage))
        .filter(|cost| !cost.is_zero())
        .unwrap_or_else(|| {
            calculate_cost_for_provider(&message.provider, &message.model_id, &message.usage)
        });
    if priced.is_zero() {
        return false;
    }
    message.cost = priced;
    true
}

#[cfg(test)]
#[path = "model_catalog_tests.rs"]
mod tests;