kcode-codex-terra-usage 0.1.1

use std::{collections::BTreeMap, fmt};

use kcode_k1_accounting::UsageValue;
use rust_decimal::Decimal;
use serde_json::Value;

type RoundUsage = [u64; 5];

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Error {
    message: String,
}

impl Error {
    fn new(message: impl Into<String>) -> Self {
        Self {
            message: message.into(),
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(&self.message)
    }
}

impl std::error::Error for Error {}

#[derive(Clone, Debug)]
pub struct UsageAccumulator {
    usage: BTreeMap<String, UsageValue>,
    total: Option<RoundUsage>,
    last: Option<RoundUsage>,
    reconciled_rounds: usize,
}

impl UsageAccumulator {
    pub fn new() -> Self {
        Self {
            usage: BTreeMap::new(),
            total: None,
            last: None,
            reconciled_rounds: 0,
        }
    }

    pub fn apply(&mut self, value: &Value) -> Result<(), Error> {
        let total = parse(
            value
                .get("total")
                .ok_or_else(|| Error::new("Codex usage omitted cumulative totals"))?,
        )?;
        let last = parse(
            value
                .get("last")
                .ok_or_else(|| Error::new("Codex usage omitted latest-round totals"))?,
        )?;

        if let Some(previous) = self.total {
            if total == previous {
                if self.last != Some(last) {
                    return Err(Error::new("Codex duplicate usage changed its latest round"));
                }
                return Ok(());
            }
            if delta(total, previous) != Some(last) {
                return Err(Error::new("Codex usage delta did not reconcile"));
            }
        } else if total != last {
            return Err(Error::new("Codex initial usage did not reconcile"));
        }

        self.add(last);
        self.total = Some(total);
        self.last = Some(last);
        self.reconciled_rounds = self.reconciled_rounds.saturating_add(1);
        Ok(())
    }

    pub fn reconciled_rounds(&self) -> usize {
        self.reconciled_rounds
    }

    pub fn snapshot(&self) -> BTreeMap<String, UsageValue> {
        self.usage.clone()
    }

    fn add(&mut self, round: RoundUsage) {
        let [input, cached, cache_write, output, reasoning] = round;
        let ordinary = input - cached - cache_write;
        let (ordinary_rate, cache_write_rate, cached_rate, output_rate) = if input > 272_000 {
            (400, 500, 40, 1_800)
        } else {
            (200, 250, 20, 1_200)
        };

        for (key, units, rate) in [
            ("input tokens", ordinary, ordinary_rate),
            ("input tokens", cache_write, cache_write_rate),
            ("cached input tokens", cached, cached_rate),
            ("output tokens", output - reasoning, output_rate),
            ("reasoning tokens", reasoning, output_rate),
        ] {
            add(&mut self.usage, key, units, rate);
        }
    }
}

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

fn delta(current: RoundUsage, previous: RoundUsage) -> Option<RoundUsage> {
    Some([
        current[0].checked_sub(previous[0])?,
        current[1].checked_sub(previous[1])?,
        current[2].checked_sub(previous[2])?,
        current[3].checked_sub(previous[3])?,
        current[4].checked_sub(previous[4])?,
    ])
}

fn parse(value: &Value) -> Result<RoundUsage, Error> {
    let part = [
        integer(value, "inputTokens")?,
        integer(value, "cachedInputTokens")?,
        integer(value, "cacheWriteInputTokens")?,
        integer(value, "outputTokens")?,
        integer(value, "reasoningOutputTokens")?,
    ];

    if part[1]
        .checked_add(part[2])
        .is_none_or(|value| value > part[0])
        || part[4] > part[3]
    {
        return Err(Error::new("Codex usage contained impossible token totals"));
    }

    Ok(part)
}

fn integer(value: &Value, key: &str) -> Result<u64, Error> {
    value
        .get(key)
        .and_then(Value::as_u64)
        .ok_or_else(|| Error::new(format!("Codex usage contained invalid {key}")))
}

fn add(usage: &mut BTreeMap<String, UsageValue>, key: &str, units: u64, rate: u64) {
    let entry = usage.entry(key.to_owned()).or_default();
    let units = Decimal::from(units);
    entry.units += units;
    entry.price_cents += units * Decimal::from(rate) / Decimal::from(1_000_000_u64);
}