ai-usagebar 1.21.1

//! Strict wire types for the official Grok Build `x.ai/billing` ACP response.

use chrono::{DateTime, Utc};
use serde::Deserialize;

use crate::error::{AppError, Result};
use crate::usage::{ResetCredits, SuperGrokPeriod, SuperGrokProduct, SuperGrokSnapshot};

const MAX_PLAN_CHARS: usize = 128;
const MAX_BENIGN_PERCENT: f64 = 100.5;
const MAX_EXACT_F64_INTEGER: i64 = 9_007_199_254_740_991;

#[derive(Debug, Clone, Deserialize, Default)]
#[serde(default)]
pub struct BillingResponse {
    pub config: Option<BillingConfig>,
    /// The ACP extension currently serializes snake_case; accept camelCase for
    /// compatibility with older/direct extension bridges.
    #[serde(alias = "subscriptionTier")]
    pub subscription_tier: Option<String>,
    /// Human-facing SKU from `/v1/settings` (e.g. "SuperGrok Heavy"). Billing's
    /// `subscription_tier` is often the short code "SuperGrok".
    #[serde(alias = "subscriptionTierDisplay")]
    pub subscription_tier_display: Option<String>,
    #[serde(skip)]
    pub reset_credits: ResetCredits,
}

#[derive(Debug, Clone, Deserialize, Default)]
#[serde(default, rename_all = "camelCase")]
pub struct BillingConfig {
    pub credit_usage_percent: Option<f64>,
    pub current_period: Option<UsagePeriod>,
    /// Deprecated legacy monthly fields returned by older Grok Build servers.
    pub monthly_limit: Option<Cent>,
    pub used: Option<Cent>,
    pub on_demand_cap: Option<Cent>,
    pub on_demand_used: Option<Cent>,
    pub prepaid_balance: Option<Cent>,
    pub is_unified_billing_user: Option<bool>,
    pub billing_period_start: Option<String>,
    pub billing_period_end: Option<String>,
    #[serde(default)]
    pub product_usage: Vec<ProductUsage>,
}

#[derive(Debug, Clone, Deserialize, Default)]
#[serde(default, rename_all = "camelCase")]
pub struct ProductUsage {
    pub product: Option<String>,
    pub usage_percent: Option<f64>,
}

#[derive(Debug, Clone, Deserialize, Default)]
#[serde(default, rename_all = "camelCase")]
pub struct UsagePeriod {
    #[serde(rename = "type")]
    pub period_type: Option<String>,
    pub start: Option<String>,
    pub end: Option<String>,
}

#[derive(Debug, Clone, Deserialize, Default)]
pub struct Cent {
    /// Proto JSON omits zero-valued scalars, so `{}` means zero. Present
    /// values must still be exact integers; fractional/saturated casts would
    /// silently corrupt billing amounts.
    #[serde(default, deserialize_with = "de_cent_val")]
    pub val: i64,
}

fn de_cent_val<'de, D>(deserializer: D) -> std::result::Result<i64, D::Error>
where
    D: serde::Deserializer<'de>,
{
    let value = serde_json::Value::deserialize(deserializer)?;
    match value {
        serde_json::Value::Number(number) => number
            .as_i64()
            .ok_or_else(|| serde::de::Error::custom("cent val must be an exact i64 integer")),
        serde_json::Value::String(text) => text
            .trim()
            .parse::<i64>()
            .map_err(|_| serde::de::Error::custom("cent val string must be an exact i64 integer")),
        _ => Err(serde::de::Error::custom(
            "cent val must be an integer number or string",
        )),
    }
}

pub fn to_snapshot(resp: BillingResponse, account_scope: &str) -> Result<SuperGrokSnapshot> {
    let plan = checked_plan(
        resp.subscription_tier_display
            .as_deref()
            .or(resp.subscription_tier.as_deref()),
    )?;
    let cfg = resp
        .config
        .ok_or_else(|| AppError::Schema("Grok Build billing response has no config".into()))?;
    let period = resolve_period(&cfg);
    let weekly_pct = resolve_usage_percent(&cfg)?;
    let reset_at = resolve_reset_at(&cfg)?;
    let products = parse_products(&cfg)?;
    let prepaid_balance = cfg
        .prepaid_balance
        .as_ref()
        .map(|cents| checked_prepaid(cents.val))
        .transpose()?;

    Ok(SuperGrokSnapshot {
        plan,
        account: account_scope.to_string(),
        weekly_pct,
        period,
        reset_at,
        prepaid_balance,
        reset_credits: resp.reset_credits,
        products,
    })
}

const MAX_PRODUCTS: usize = 16;

fn parse_products(cfg: &BillingConfig) -> Result<Vec<SuperGrokProduct>> {
    let mut products = Vec::new();
    for item in cfg.product_usage.iter().take(MAX_PRODUCTS) {
        let Some(raw) = item
            .product
            .as_deref()
            .map(str::trim)
            .filter(|s| !s.is_empty())
        else {
            continue;
        };
        if raw.chars().count() > MAX_PLAN_CHARS || raw.chars().any(char::is_control) {
            return Err(AppError::Schema(
                "Grok Build product name is invalid".into(),
            ));
        }
        let percent = match item.usage_percent {
            Some(value) => checked_percent(value)?,
            None => 0,
        };
        products.push(SuperGrokProduct {
            label: product_label(raw),
            percent,
        });
    }
    Ok(products)
}

fn product_label(raw: &str) -> String {
    match raw {
        "GrokBuild" => "Grok Build".into(),
        "GrokChat" => "Grok Chat".into(),
        "GrokImagine" => "Grok Imagine".into(),
        "GrokTasks" => "Grok Tasks".into(),
        "Api" => "xAI API".into(),
        other if other.starts_with("Grok") && other.len() > 4 => {
            let rest = &other[4..];
            let spaced = rest.chars().fold(String::new(), |mut out, ch| {
                if ch.is_uppercase() && !out.is_empty() {
                    out.push(' ');
                }
                out.push(ch);
                out
            });
            if spaced.is_empty() {
                "Grok".into()
            } else {
                format!("Grok {spaced}")
            }
        }
        other => other.to_string(),
    }
}

fn checked_plan(value: Option<&str>) -> Result<String> {
    let value = value.map(str::trim).filter(|s| !s.is_empty());
    let Some(value) = value else {
        return Ok("SuperGrok".into());
    };
    if value.chars().count() > MAX_PLAN_CHARS || value.chars().any(char::is_control) {
        return Err(AppError::Schema(
            "Grok Build subscription tier is invalid".into(),
        ));
    }
    Ok(value.to_string())
}

fn resolve_usage_percent(cfg: &BillingConfig) -> Result<i32> {
    if let Some(percent) = cfg.credit_usage_percent {
        return checked_percent(percent);
    }

    // Proto JSON omits the zero-valued percentage immediately after rollover.
    // A typed current period makes that omission unambiguous. Never splice in
    // deprecated monthly counters under a weekly current-period reset.
    if cfg.current_period.is_some() {
        return Ok(0);
    }

    match (&cfg.used, &cfg.monthly_limit) {
        (Some(used), Some(limit)) if limit.val > 0 && used.val >= 0 => {
            checked_percent((used.val as f64 / limit.val as f64) * 100.0)
        }
        (Some(_), Some(_)) => Err(AppError::Schema(
            "Grok Build legacy billing counters are negative or have a non-positive limit".into(),
        )),
        _ if cfg.billing_period_end.is_some()
            || cfg.prepaid_balance.is_some()
            || cfg.is_unified_billing_user.is_some() =>
        {
            Ok(0)
        }
        _ => Err(AppError::Schema(
            "Grok Build billing response has no usage percentage or coherent legacy counters"
                .into(),
        )),
    }
}

fn checked_percent(value: f64) -> Result<i32> {
    if !value.is_finite() || !(0.0..=MAX_BENIGN_PERCENT).contains(&value) {
        return Err(AppError::Schema(
            "Grok Build billing percentage is outside the supported range".into(),
        ));
    }
    Ok(value.round().clamp(0.0, 100.0) as i32)
}

fn resolve_period(cfg: &BillingConfig) -> SuperGrokPeriod {
    let raw = cfg
        .current_period
        .as_ref()
        .and_then(|period| period.period_type.as_deref())
        .unwrap_or_default();
    if raw.ends_with("WEEKLY") {
        SuperGrokPeriod::Weekly
    } else if raw.ends_with("MONTHLY")
        || (cfg.current_period.is_none()
            && (cfg.monthly_limit.is_some()
                || cfg.used.is_some()
                || cfg.billing_period_end.is_some()))
    {
        SuperGrokPeriod::Monthly
    } else {
        SuperGrokPeriod::Unknown
    }
}

fn resolve_reset_at(cfg: &BillingConfig) -> Result<Option<DateTime<Utc>>> {
    if let Some(period) = cfg.current_period.as_ref() {
        return parse_optional_datetime(period.end.as_deref(), "currentPeriod.end");
    }
    parse_optional_datetime(cfg.billing_period_end.as_deref(), "billingPeriodEnd")
}

fn parse_optional_datetime(value: Option<&str>, field: &str) -> Result<Option<DateTime<Utc>>> {
    let Some(value) = value else {
        return Ok(None);
    };
    DateTime::parse_from_rfc3339(value)
        .map(|dt| Some(dt.with_timezone(&Utc)))
        .map_err(|_| AppError::Schema(format!("Grok Build {field} is not RFC 3339")))
}

fn checked_prepaid(cents: i64) -> Result<f64> {
    if !(0..=MAX_EXACT_F64_INTEGER).contains(&cents) {
        return Err(AppError::Schema(
            "Grok Build prepaid balance is negative or too large to represent exactly".into(),
        ));
    }
    Ok(cents as f64 / 100.0)
}

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

    #[test]
    fn weekly_acp_shape_is_coherent() {
        let response: BillingResponse = serde_json::from_str(
            r#"{
              "config": {
                "creditUsagePercent": 42.5,
                "currentPeriod": {
                  "type": "USAGE_PERIOD_TYPE_WEEKLY",
                  "end": "2026-08-10T00:00:00Z"
                },
                "prepaidBalance": {"val": 1250}
              },
              "subscription_tier": "SuperGrok Heavy"
            }"#,
        )
        .unwrap();
        let snapshot = to_snapshot(response, "opaque-scope").unwrap();
        assert_eq!(snapshot.weekly_pct, 43);
        assert_eq!(snapshot.period, SuperGrokPeriod::Weekly);
        assert_eq!(snapshot.plan, "SuperGrok Heavy");
        assert_eq!(snapshot.prepaid_balance, Some(12.5));
        assert!(snapshot.products.is_empty());
    }

    #[test]
    fn product_usage_rows_are_labelled_and_rounded() {
        let response: BillingResponse = serde_json::from_str(
            r#"{
              "config": {
                "creditUsagePercent": 90.0,
                "currentPeriod": {
                  "type": "USAGE_PERIOD_TYPE_WEEKLY",
                  "end": "2026-09-20T13:26:44Z"
                },
                "productUsage": [
                  {"product": "GrokBuild", "usagePercent": 87.4},
                  {"product": "GrokChat", "usagePercent": 2.6},
                  {"product": "GrokImagine"},
                  {"product": "Api", "usagePercent": 0.0}
                ]
              }
            }"#,
        )
        .unwrap();
        let snapshot = to_snapshot(response, "scope").unwrap();
        assert_eq!(
            snapshot
                .products
                .iter()
                .map(|p| (p.label.as_str(), p.percent))
                .collect::<Vec<_>>(),
            vec![
                ("Grok Build", 87),
                ("Grok Chat", 3),
                ("Grok Imagine", 0),
                ("xAI API", 0),
            ]
        );
    }

    #[test]
    fn subscription_tier_display_wins_over_the_short_code() {
        let response: BillingResponse = serde_json::from_str(
            r#"{
              "config": {
                "creditUsagePercent": 10,
                "currentPeriod": { "type": "USAGE_PERIOD_TYPE_WEEKLY" }
              },
              "subscription_tier": "SuperGrok",
              "subscription_tier_display": "SuperGrok Heavy"
            }"#,
        )
        .unwrap();
        let snapshot = to_snapshot(response, "scope").unwrap();
        assert_eq!(snapshot.plan, "SuperGrok Heavy");
    }

    #[test]
    fn legacy_monthly_shape_keeps_its_own_reset() {
        let response: BillingResponse = serde_json::from_str(
            r#"{"config":{"monthlyLimit":{"val":"2000"},"used":{"val":500},"billingPeriodEnd":"2026-09-01T00:00:00Z"}}"#,
        )
        .unwrap();
        let snapshot = to_snapshot(response, "scope").unwrap();
        assert_eq!(snapshot.weekly_pct, 25);
        assert_eq!(snapshot.period, SuperGrokPeriod::Monthly);
        assert_eq!(
            snapshot.reset_at.unwrap().to_rfc3339(),
            "2026-09-01T00:00:00+00:00"
        );
    }

    #[test]
    fn omitted_zero_percent_does_not_import_legacy_monthly_usage() {
        let response: BillingResponse = serde_json::from_str(
            r#"{"config":{"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","end":"2026-08-13T00:00:00Z"},"monthlyLimit":{"val":1000},"used":{"val":900}}}"#,
        )
        .unwrap();
        let snapshot = to_snapshot(response, "scope").unwrap();
        assert_eq!(snapshot.weekly_pct, 0);
        assert_eq!(snapshot.period, SuperGrokPeriod::Weekly);
    }

    #[test]
    fn cents_must_be_exact_integers() {
        for value in ["1.5", "1e100", "null", "true"] {
            let body = format!(r#"{{"config":{{"prepaidBalance":{{"val":{value}}}}}}}"#);
            assert!(
                serde_json::from_str::<BillingResponse>(&body).is_err(),
                "{body}"
            );
        }
        let omitted: BillingResponse =
            serde_json::from_str(r#"{"config":{"prepaidBalance":{}}}"#).unwrap();
        assert_eq!(omitted.config.unwrap().prepaid_balance.unwrap().val, 0);
    }

    #[test]
    fn malformed_percentages_and_resets_are_rejected() {
        for percent in [-1.0, 101.0, f64::INFINITY] {
            assert!(checked_percent(percent).is_err());
        }
        let response: BillingResponse = serde_json::from_str(
            r#"{"config":{"creditUsagePercent":5,"currentPeriod":{"end":"not-a-date"}}}"#,
        )
        .unwrap();
        assert!(to_snapshot(response, "scope").is_err());
        assert!(checked_prepaid(-1).is_err());
        assert!(checked_prepaid(MAX_EXACT_F64_INTEGER + 1).is_err());
    }

    #[test]
    fn plan_labels_are_bounded_and_control_free() {
        assert!(checked_plan(Some(&"x".repeat(MAX_PLAN_CHARS + 1))).is_err());
        assert!(checked_plan(Some("bad\u{1b}[31m")).is_err());
        assert_eq!(checked_plan(Some("  ")).unwrap(), "SuperGrok");
    }
}