o402 0.1.5

//! Integer atomic ceiling and meter. Human strings parse only at load.

use std::str::FromStr as _;

use r402_protocol::MoneyAmount;
use rust_decimal::Decimal;

use super::usage::Usage;
use crate::config::{ConfigError, ModelConfig, PricingConfig, PricingDefault, UsagePolicy};

/// Atomic units per million tokens.
const MILLION: u128 = 1_000_000;

/// Model scheme after inheriting `[pricing.default]`.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum Scheme {
    /// Per-request `price`.
    Exact,
    /// EVM Permit2 authorization; settle `actual ≤ ceiling`.
    Upto,
}

impl Scheme {
    /// Wire scheme name.
    #[must_use]
    pub(crate) const fn as_str(self) -> &'static str {
        match self {
            Self::Exact => "exact",
            Self::Upto => "upto",
        }
    }
}

/// Per-million token rates in atomic units.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) struct TokenPpm {
    /// Uncached input.
    pub(crate) input: u128,
    /// Non-reasoning output. Unused for embeddings.
    pub(crate) output: u128,
    /// Cached input.
    pub(crate) cached: u128,
    /// Reasoning output.
    pub(crate) reasoning: u128,
}

/// `ceiling_multiplier` as an integer fraction ≥ 1.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) struct Multiplier {
    /// Numerator.
    numer: u128,
    /// Denominator.
    denom: u128,
}

impl Multiplier {
    const IDENTITY: Self = Self { numer: 1, denom: 1 };

    const fn apply(self, value: u128) -> u128 {
        if self.denom == 0 || (self.numer == 1 && self.denom == 1) {
            return value;
        }
        value.saturating_mul(self.numer) / self.denom
    }
}

/// Rates converted at config load. Request path does not parse money strings.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) struct LoadedRates {
    /// Model scheme.
    pub(crate) scheme: Scheme,
    /// Exact per-request atomic price. Zero on `upto`.
    pub(crate) exact_price: u128,
    /// Request floor.
    pub(crate) floor: u128,
    /// Hard ceiling cap.
    pub(crate) max_ceiling: u128,
    /// Max input tokens for the 402 ceiling.
    pub(crate) max_input_tokens: u32,
    /// Fallback max output tokens when the request omits a cap.
    pub(crate) default_max_output: u32,
    /// Atomic-per-million token rates.
    pub(crate) ppm: TokenPpm,
    /// Ceiling multiplier.
    multiplier: Multiplier,
}

/// Parses a model's inherited rates into atomic integers.
///
/// # Errors
///
/// Missing required fields, unparsable money strings, or `ceiling_multiplier` < 1.
pub(crate) fn load_rates(
    model: &ModelConfig,
    pricing: Option<&PricingConfig>,
    decimals: u8,
) -> Result<LoadedRates, ConfigError> {
    let default = pricing.and_then(|pricing| pricing.default.as_ref());
    let scheme = parse_scheme(inherit_str(model.scheme.as_deref(), default, |row| {
        row.scheme.as_deref()
    }))?;
    match scheme {
        Scheme::Exact => load_exact(model, default, decimals),
        Scheme::Upto => load_upto(model, default, decimals),
    }
}

/// Exact rates from inherited `price`.
///
/// # Errors
///
/// Missing or unparsable `price`.
pub(crate) fn load_exact_rates(
    model: &ModelConfig,
    pricing: Option<&PricingConfig>,
    decimals: u8,
) -> Result<LoadedRates, ConfigError> {
    let default = pricing.and_then(|pricing| pricing.default.as_ref());
    load_exact(model, default, decimals)
}

/// Upto rates from inherited per-million fields.
///
/// # Errors
///
/// Missing required upto fields or unparsable amounts.
pub(crate) fn load_upto_rates(
    model: &ModelConfig,
    pricing: Option<&PricingConfig>,
    decimals: u8,
) -> Result<LoadedRates, ConfigError> {
    let default = pricing.and_then(|pricing| pricing.default.as_ref());
    load_upto(model, default, decimals)
}

/// 402 amount: exact price, or the integer upto ceiling.
#[must_use]
pub(crate) fn ceiling_or_price(rates: &LoadedRates, max_out: u128, embeddings: bool) -> u128 {
    match rates.scheme {
        Scheme::Exact => rates.exact_price,
        Scheme::Upto => upto_ceiling(rates, max_out, embeddings),
    }
}

/// Output cap: request `max_*` if set, else the model's default.
#[must_use]
pub(crate) fn max_out(
    rates: &LoadedRates,
    max_tokens: Option<u32>,
    max_completion_tokens: Option<u32>,
    max_output_tokens: Option<u32>,
) -> u128 {
    u128::from(
        max_completion_tokens
            .or(max_output_tokens)
            .or(max_tokens)
            .unwrap_or(rates.default_max_output),
    )
}

/// Metered atomic amount from usage. Cached/reasoning are subtracted first.
#[must_use]
pub(crate) fn meter(rates: &LoadedRates, usage: &Usage) -> u128 {
    let mut total = rates.floor;
    total = total.saturating_add(token_term(usage.uncached(), rates.ppm.input));
    total = total.saturating_add(token_term(usage.cached(), rates.ppm.cached));
    total = total.saturating_add(token_term(usage.output(), rates.ppm.output));
    total.saturating_add(token_term(usage.reasoning(), rates.ppm.reasoning))
}

/// Settle amount: exact tag, metered `min(meter, ceiling)`, or missing-usage policy.
#[must_use]
pub(crate) fn actual(
    rates: &LoadedRates,
    ceiling: u128,
    usage: Option<&Usage>,
    missing: UsagePolicy,
) -> u128 {
    match rates.scheme {
        Scheme::Exact => ceiling,
        Scheme::Upto => usage.map_or_else(
            || match missing {
                UsagePolicy::Ceiling => ceiling,
                UsagePolicy::Floor => rates.floor.min(ceiling),
            },
            |usage| meter(rates, usage).min(ceiling),
        ),
    }
}

fn upto_ceiling(rates: &LoadedRates, max_out: u128, embeddings: bool) -> u128 {
    let input_term = token_term(u128::from(rates.max_input_tokens), rates.ppm.input);
    let output_term = if embeddings {
        0
    } else {
        token_term(max_out, rates.ppm.output)
    };
    let raw = input_term
        .saturating_add(output_term)
        .saturating_add(rates.floor);
    let scaled = rates.multiplier.apply(raw);
    scaled.max(rates.floor).min(rates.max_ceiling)
}

const fn token_term(tokens: u128, ppm: u128) -> u128 {
    tokens.saturating_mul(ppm) / MILLION
}

fn load_exact(
    model: &ModelConfig,
    default: Option<&PricingDefault>,
    decimals: u8,
) -> Result<LoadedRates, ConfigError> {
    let price = inherit_str(model.price.as_deref(), default, |row| row.price.as_deref())
        .ok_or_else(|| {
            ConfigError::Validation(format!(
                "model '{}' requires an exact price when payment is enabled",
                model.id
            ))
        })?;
    Ok(LoadedRates {
        scheme: Scheme::Exact,
        exact_price: exact_atomic(price, decimals)?,
        floor: 0,
        max_ceiling: 0,
        max_input_tokens: 0,
        default_max_output: 0,
        ppm: TokenPpm {
            input: 0,
            output: 0,
            cached: 0,
            reasoning: 0,
        },
        multiplier: Multiplier::IDENTITY,
    })
}

fn load_upto(
    model: &ModelConfig,
    default: Option<&PricingDefault>,
    decimals: u8,
) -> Result<LoadedRates, ConfigError> {
    let input = inherit_str(model.input_per_million.as_deref(), default, |row| {
        row.input_per_million.as_deref()
    })
    .ok_or_else(|| {
        ConfigError::Validation(format!(
            "model '{}' requires input_per_million when scheme is upto",
            model.id
        ))
    })?;
    let max_input_tokens = inherit_u32(model.max_input_tokens, default, |row| row.max_input_tokens)
        .ok_or_else(|| {
            ConfigError::Validation(format!(
                "model '{}' requires max_input_tokens when scheme is upto",
                model.id
            ))
        })?;
    let output = inherit_str(model.output_per_million.as_deref(), default, |row| {
        row.output_per_million.as_deref()
    })
    .ok_or_else(|| {
        ConfigError::Validation(format!(
            "model '{}' requires output_per_million when scheme is upto",
            model.id
        ))
    })?;
    let cached = inherit_str(model.cached_input_per_million.as_deref(), default, |row| {
        row.cached_input_per_million.as_deref()
    });
    let reasoning = inherit_str(model.reasoning_per_million.as_deref(), default, |row| {
        row.reasoning_per_million.as_deref()
    });
    let floor_raw = inherit_str(model.request_floor.as_deref(), default, |row| {
        row.request_floor.as_deref()
    });
    let max_ceiling_raw = inherit_str(model.max_ceiling.as_deref(), default, |row| {
        row.max_ceiling.as_deref()
    });
    let multiplier = inherit_str(model.ceiling_multiplier.as_deref(), default, |row| {
        row.ceiling_multiplier.as_deref()
    });
    let floor = optional_atomic(floor_raw, decimals, "request_floor")?;
    let max_ceiling = match max_ceiling_raw {
        Some(raw) => atomic_amount(raw, decimals, "max_ceiling")?,
        None => u128::MAX,
    };
    if max_ceiling < floor {
        return Err(ConfigError::Validation(format!(
            "model '{}' max_ceiling is below request_floor",
            model.id
        )));
    }
    Ok(LoadedRates {
        scheme: Scheme::Upto,
        exact_price: 0,
        floor,
        max_ceiling,
        max_input_tokens,
        default_max_output: inherit_u32(model.default_max_output_tokens, default, |row| {
            row.default_max_output_tokens
        })
        .unwrap_or(0),
        ppm: TokenPpm {
            input: atomic_amount(input, decimals, "input_per_million")?,
            output: atomic_amount(output, decimals, "output_per_million")?,
            cached: optional_atomic(cached, decimals, "cached_input_per_million")?,
            reasoning: optional_atomic(reasoning, decimals, "reasoning_per_million")?,
        },
        multiplier: parse_multiplier(multiplier)?,
    })
}

fn parse_scheme(raw: Option<&str>) -> Result<Scheme, ConfigError> {
    match raw {
        Some("exact") => Ok(Scheme::Exact),
        Some("upto") => Ok(Scheme::Upto),
        Some(other) => Err(ConfigError::Validation(format!(
            "scheme '{other}' is not supported"
        ))),
        None => Err(ConfigError::Validation(
            "scheme or pricing.default.scheme is required when payment is enabled".to_owned(),
        )),
    }
}

fn parse_multiplier(raw: Option<&str>) -> Result<Multiplier, ConfigError> {
    let Some(raw) = raw else {
        return Ok(Multiplier::IDENTITY);
    };
    let value = Decimal::from_str(raw).map_err(|error| {
        ConfigError::Validation(format!("invalid ceiling_multiplier '{raw}': {error}"))
    })?;
    if value.is_sign_negative() || value < Decimal::ONE {
        return Err(ConfigError::Validation(format!(
            "ceiling_multiplier must be >= 1 (got '{raw}')"
        )));
    }
    let scale = value.scale();
    let denom = 10u128.checked_pow(scale).ok_or_else(|| {
        ConfigError::Validation(format!("ceiling_multiplier '{raw}' scale is out of range"))
    })?;
    let numer = value.mantissa().unsigned_abs();
    if denom == 0 {
        return Ok(Multiplier::IDENTITY);
    }
    Ok(Multiplier { numer, denom })
}

fn optional_atomic(raw: Option<&str>, decimals: u8, what: &str) -> Result<u128, ConfigError> {
    raw.map_or(Ok(0), |raw| atomic_amount(raw, decimals, what))
}

/// Parses a human-readable token amount into atomic units.
///
/// # Errors
///
/// Returns [`ConfigError::Validation`] when `price` is not a non-negative
/// decimal that fits `decimals`.
pub(crate) fn exact_atomic(price: &str, decimals: u8) -> Result<u128, ConfigError> {
    atomic_amount(price, decimals, "exact price")
}

fn atomic_amount(raw: &str, decimals: u8, what: &str) -> Result<u128, ConfigError> {
    if raw
        .trim()
        .parse::<Decimal>()
        .is_ok_and(|value| value == Decimal::ZERO)
    {
        return Ok(0);
    }
    MoneyAmount::parse(raw)
        .and_then(|amount| amount.to_token_amount(decimals))
        .map_err(|error| ConfigError::Validation(format!("invalid {what} '{raw}': {error}")))
}

fn inherit_str<'a>(
    model: Option<&'a str>,
    default: Option<&'a PricingDefault>,
    pick: impl Fn(&'a PricingDefault) -> Option<&'a str>,
) -> Option<&'a str> {
    model.or_else(|| default.and_then(pick))
}

fn inherit_u32(
    model: Option<u32>,
    default: Option<&PricingDefault>,
    pick: impl Fn(&PricingDefault) -> Option<u32>,
) -> Option<u32> {
    model.or_else(|| default.and_then(pick))
}

#[cfg(test)]
mod tests {
    use super::{
        LoadedRates, MILLION, Multiplier, Scheme, TokenPpm, actual, ceiling_or_price, exact_atomic,
        meter,
    };
    use crate::config::UsagePolicy;
    use crate::payment::usage::Usage;

    fn mini() -> LoadedRates {
        LoadedRates {
            scheme: Scheme::Upto,
            exact_price: 0,
            floor: 10,
            max_ceiling: 5_000_000,
            max_input_tokens: 128_000,
            default_max_output: 16_384,
            ppm: TokenPpm {
                input: 150_000,
                output: 600_000,
                cached: 75_000,
                reasoning: 600_000,
            },
            multiplier: Multiplier::IDENTITY,
        }
    }

    fn embedding() -> LoadedRates {
        LoadedRates {
            scheme: Scheme::Upto,
            exact_price: 0,
            floor: 10,
            max_ceiling: 5_000_000,
            max_input_tokens: 8191,
            default_max_output: 16_384,
            ppm: TokenPpm {
                input: 20_000,
                output: 600_000,
                cached: 0,
                reasoning: 0,
            },
            multiplier: Multiplier::IDENTITY,
        }
    }

    #[test]
    fn zero_rate_is_zero_atomic() {
        assert_eq!(exact_atomic("0", 6).expect("zero"), 0, "0");
        assert_eq!(exact_atomic("0.0", 6).expect("zero"), 0, "0.0");
    }

    #[test]
    fn usdc_milli_is_one_thousand_atomic() {
        assert_eq!(
            exact_atomic("0.001", 6).expect("parse"),
            1000,
            "0.001 USDC-6"
        );
    }

    #[test]
    fn rejects_more_fractional_digits_than_token() {
        let err = exact_atomic("0.0000001", 6).expect_err("precision");
        assert!(err.to_string().contains("invalid exact price"), "got {err}");
    }

    #[test]
    fn gpt4o_mini_ceiling_is_29040() {
        let rates = mini();
        assert_eq!(128_000 * 150_000 / MILLION, 19_200, "input_term");
        assert_eq!(16_384 * 600_000 / MILLION, 9_830, "output_term");
        assert_eq!(
            ceiling_or_price(&rates, u128::from(rates.default_max_output), false),
            29_040,
            "USDC-6 gpt-4o-mini"
        );
    }

    #[test]
    fn cached_and_reasoning_are_not_double_counted() {
        let usage = Usage {
            prompt_tokens: 1000,
            completion_tokens: 100,
            cached_tokens: 200,
            reasoning_tokens: 40,
        };
        assert_eq!(usage.uncached(), 800, "uncached");
        assert_eq!(usage.output(), 60, "output");
        let billed = meter(&mini(), &usage);
        assert_eq!(billed, 205, "floor+uncached+cached+output+reasoning");
        let double = mini().floor
            + 1000 * 150_000 / MILLION
            + 200 * 75_000 / MILLION
            + 100 * 600_000 / MILLION
            + 40 * 600_000 / MILLION;
        assert!(billed < double, "{billed} vs double-count {double}");
    }

    #[test]
    fn meter_clamps_to_ceiling() {
        let rates = mini();
        let ceiling = ceiling_or_price(&rates, u128::from(rates.default_max_output), false);
        let usage = Usage {
            prompt_tokens: 128_000,
            completion_tokens: 20_000,
            cached_tokens: 0,
            reasoning_tokens: 0,
        };
        let metered = meter(&rates, &usage);
        assert!(metered > ceiling, "meter {metered} > ceiling {ceiling}");
        assert_eq!(
            actual(&rates, ceiling, Some(&usage), UsagePolicy::Ceiling),
            ceiling,
            "clamp"
        );
    }

    #[test]
    fn embeddings_ceiling_has_no_output_term() {
        let rates = embedding();
        assert_eq!(8191 * 20_000 / MILLION, 163, "input_term");
        assert_eq!(
            ceiling_or_price(&rates, u128::from(rates.default_max_output), true),
            173,
            "embeddings ceiling"
        );
        let with_output = ceiling_or_price(&rates, u128::from(rates.default_max_output), false);
        assert!(
            with_output > 173,
            "output term must not apply to embeddings: {with_output}"
        );
    }

    #[test]
    fn missing_usage_follows_policy() {
        let rates = mini();
        let ceiling = 29_040;
        assert_eq!(
            actual(&rates, ceiling, None, UsagePolicy::Ceiling),
            ceiling,
            "ceiling policy"
        );
        assert_eq!(
            actual(&rates, ceiling, None, UsagePolicy::Floor),
            10,
            "floor policy"
        );
    }

    #[test]
    fn missing_usage_floor_does_not_exceed_ceiling() {
        let mut rates = mini();
        rates.floor = 100;
        assert_eq!(
            actual(&rates, 40, None, UsagePolicy::Floor),
            40,
            "floor clamped to ceiling"
        );
    }
}