obol-core 0.2.2

//! obol-core: parse agent transcripts and estimate token cost.

pub mod cost;
pub mod error;
pub mod model;
pub mod pricing;
pub mod transcript;

pub use error::ObolError;
pub use model::{
    Approximation, CostEstimate, MessageUsage, ModelCost, PricingSource, Provider, TokenBuckets,
};
pub use transcript::Dialect;

use std::path::{Path, PathBuf};

/// Report from a pricing refresh.
#[derive(Debug, serde::Serialize)]
pub struct RefreshReport {
    pub models: usize,
    pub as_of: String,
    pub written_to: PathBuf,
}

/// Resolve the price snapshot. Explicit OBOL_PRICING_DIR wins absolutely; otherwise
/// pick whichever of {on-disk current.json, embedded} has the newer `as_of`, on-disk
/// winning ties; embedded is the floor.
fn resolve_store() -> Result<(pricing::PriceStore, PricingSource), ObolError> {
    if std::env::var_os("OBOL_PRICING_DIR").is_some() {
        let store = pricing::PriceStore::load(&pricing::current_path())?;
        return Ok((store, PricingSource::Local));
    }
    let embedded = pricing::embedded()?;
    let local_path = pricing::current_path();
    if local_path.exists() {
        if let Ok(local) = pricing::PriceStore::load(&local_path) {
            if local.as_of >= embedded.as_of {
                return Ok((local, PricingSource::Local));
            }
        }
    }
    Ok((embedded, PricingSource::Bundled))
}

/// Estimate the cost of a transcript file under the given dialect. Loads the active
/// price snapshot (bundled fallback) and prices the parsed usage.
pub fn estimate_cost(path: &Path, dialect: Dialect) -> Result<CostEstimate, ObolError> {
    let (store, source_kind) = resolve_store()?;
    let bytes = std::fs::read(path)?;
    let usages = transcript::parse(&bytes, dialect)?;
    Ok(cost::estimate(&usages, &store, source_kind))
}

/// Fetch the LiteLLM sheet and write it as the active snapshot. `as_of` is the
/// caller's date string (the library has no clock).
pub fn refresh_pricing_tables(as_of: &str) -> Result<RefreshReport, ObolError> {
    let mut store = pricing::refresh::fetch_litellm(as_of)?; // {litellm: …}
    let openrouter = pricing::refresh::fetch_openrouter()?;
    store
        .namespaces
        .insert("openrouter".to_string(), openrouter);
    let models: usize = store.namespaces.values().map(|m| m.len()).sum();
    let dir = pricing::pricing_dir();
    store.save(&dir.join(format!("prices-{as_of}.json")))?;
    let current = pricing::current_path();
    store.save(&current)?;
    Ok(RefreshReport {
        models,
        as_of: as_of.to_string(),
        written_to: current,
    })
}

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

    #[test]
    fn estimate_cost_on_bytes_with_missing_tables_errors() {
        std::env::set_var("OBOL_PRICING_DIR", "/nonexistent/obol-xyz");
        let tmp = std::env::temp_dir().join(format!("obol-t-{}-{}", std::process::id(), line!()));
        std::fs::write(
            &tmp,
            include_bytes!("../tests/fixtures/claude-mini.jsonl").as_slice(),
        )
        .unwrap();
        assert!(matches!(
            estimate_cost(&tmp, Dialect::Claude),
            Err(ObolError::PricingTablesMissing(_))
        ));
        std::fs::remove_file(&tmp).ok();
        std::env::remove_var("OBOL_PRICING_DIR");
    }

    #[test]
    fn estimate_cost_end_to_end_with_seeded_store() {
        let dir = std::env::temp_dir().join(format!("obol-api-{}", std::process::id()));
        std::env::set_var("OBOL_PRICING_DIR", &dir);
        // seed the store from the sample sheet
        let store = pricing::refresh::normalize_litellm(
            include_bytes!("../tests/fixtures/litellm-sample.json"),
            "2026-06-04",
        )
        .unwrap();
        store.save(&pricing::current_path()).unwrap();

        let tmp = std::env::temp_dir().join(format!("obol-t-{}-{}", std::process::id(), line!()));
        std::fs::write(
            &tmp,
            include_bytes!("../tests/fixtures/claude-mini.jsonl").as_slice(),
        )
        .unwrap();
        let est = estimate_cost(&tmp, Dialect::Claude).unwrap();
        assert!(est.total_usd > 0.0);
        assert_eq!(est.pricing_as_of, "2026-06-04");
        std::fs::remove_file(&tmp).ok();

        std::fs::remove_dir_all(&dir).ok();
        std::env::remove_var("OBOL_PRICING_DIR");
    }

    #[test]
    fn estimate_cost_from_path_then_detect() {
        let dir = std::env::temp_dir().join(format!("obol-path-{}", std::process::id()));
        std::env::set_var("OBOL_PRICING_DIR", &dir);
        let store = pricing::refresh::normalize_litellm(
            include_bytes!("../tests/fixtures/litellm-sample.json"),
            "2026-06-04",
        )
        .unwrap();
        store.save(&pricing::current_path()).unwrap();

        // Write the Claude fixture to a real file, detect dialect, then price.
        let transcript = dir.join("session.jsonl");
        std::fs::write(
            &transcript,
            include_bytes!("../tests/fixtures/claude-mini.jsonl"),
        )
        .unwrap();
        let bytes = std::fs::read(&transcript).unwrap();
        let d = transcript::detect(&bytes).unwrap();
        let est = estimate_cost(&transcript, d).unwrap();
        assert!(est.total_usd > 0.0);

        std::fs::remove_dir_all(&dir).ok();
        std::env::remove_var("OBOL_PRICING_DIR");
    }

    #[test]
    fn falls_back_to_embedded_when_no_local_snapshot() {
        // Force "no on-disk snapshot" hermetically: point XDG at an empty dir so
        // `current_path()` resolves to a nonexistent file and the embedded snapshot
        // is used. (Setting OBOL_PRICING_DIR instead would take the explicit-override
        // branch and error PricingTablesMissing rather than fall back to embedded.)
        let xdg = std::env::temp_dir().join(format!("obol-xdg-{}", std::process::id()));
        std::fs::create_dir_all(&xdg).unwrap();
        std::env::remove_var("OBOL_PRICING_DIR");
        std::env::set_var("XDG_DATA_HOME", &xdg);
        let tmp = std::env::temp_dir().join(format!("obol-t-{}-{}", std::process::id(), line!()));
        std::fs::write(
            &tmp,
            include_bytes!("../tests/fixtures/claude-mini.jsonl").as_slice(),
        )
        .unwrap();
        let est = estimate_cost(&tmp, Dialect::Claude).unwrap();
        assert_eq!(est.pricing_source, crate::model::PricingSource::Bundled);
        assert!(est.total_usd > 0.0, "embedded snapshot should price claude");
        std::fs::remove_file(&tmp).ok();
        std::env::remove_var("XDG_DATA_HOME");
        std::fs::remove_dir_all(&xdg).ok();
    }

    #[test]
    fn explicit_override_uses_local_source() {
        let dir = std::env::temp_dir().join(format!("obol-resolve-{}", std::process::id()));
        std::env::set_var("OBOL_PRICING_DIR", &dir);
        let store = pricing::refresh::normalize_litellm(
            include_bytes!("../tests/fixtures/litellm-sample.json"),
            "2099-01-01",
        )
        .unwrap();
        store.save(&pricing::current_path()).unwrap();
        let tmp = std::env::temp_dir().join(format!("obol-t-{}-{}", std::process::id(), line!()));
        std::fs::write(
            &tmp,
            include_bytes!("../tests/fixtures/claude-mini.jsonl").as_slice(),
        )
        .unwrap();
        let est = estimate_cost(&tmp, Dialect::Claude).unwrap();
        assert_eq!(est.pricing_source, crate::model::PricingSource::Local);
        std::fs::remove_file(&tmp).ok();
        std::fs::remove_dir_all(&dir).ok();
        std::env::remove_var("OBOL_PRICING_DIR");
    }

    #[test]
    fn kimi_model_surfaces_unpriced_loudly() {
        std::env::remove_var("OBOL_PRICING_DIR");
        let tmp = std::env::temp_dir().join(format!("obol-t-{}-{}", std::process::id(), line!()));
        std::fs::write(
            &tmp,
            include_bytes!("../tests/fixtures/kimi-mini.jsonl").as_slice(),
        )
        .unwrap();
        let est = estimate_cost(&tmp, Dialect::Kimi).unwrap();
        assert_eq!(est.total_usd, 0.0, "kimi-for-coding is unpriced -> $0");
        assert!(
            est.unpriced_models.contains(&"kimi-for-coding".to_string()),
            "must name the unpriced model: {:?}",
            est.unpriced_models
        );
        std::fs::remove_file(&tmp).ok();
    }

    #[test]
    fn refresh_report_serializes() {
        let r = RefreshReport {
            models: 7,
            as_of: "2026-06-05".into(),
            written_to: "/x/current.json".into(),
        };
        let v = serde_json::to_value(&r).unwrap();
        assert_eq!(v["models"], 7);
        assert_eq!(v["as_of"], "2026-06-05");
        assert_eq!(v["written_to"], "/x/current.json");
    }
}