cortex-rs-stats 0.2.0

// relay-stats: usage + cost dashboard
// Purpose: fetch rolling 7-day token/cost data from provider APIs or local log scraper
// Public surface: StatsEngine (new, fetch_summary)
// NOT responsible for: display/formatting (handled by relay-cli), memory, routing
// Test strategy: unit test log-scraper against fixture JSONL; API path requires integration test;
//                README code example is compile-tested via doctest below.

#![doc = include_str!("../README.md")]

mod log_scraper;
mod provider_api;
mod savings;

pub use log_scraper::{scrape_claude_logs, TokenCounts};

use anyhow::Result;
use cortex_rs_core::{Db, CortexConfig};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

pub struct StatsEngine {
    config: CortexConfig,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UsageSummary {
    pub days: u32,
    pub by_model: Vec<ModelUsage>,
    pub total_cost_usd: f64,
    pub routing_savings_usd: f64,
    pub source: StatsSource,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelUsage {
    pub model: String,
    pub input_tokens: u64,
    pub output_tokens: u64,
    #[serde(default)]
    pub cache_creation_input_tokens: u64,
    #[serde(default)]
    pub cache_read_input_tokens: u64,
    pub cost_usd: f64,
    pub pct: f64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum StatsSource {
    /// Fetched from provider org API with admin key
    ProviderApi,
    /// Parsed from ~/.claude/ session JSONL files
    LocalLogScraper,
    /// No data available
    Unavailable,
}

impl StatsEngine {
    pub fn new(config: CortexConfig) -> Self {
        StatsEngine { config }
    }

    pub async fn fetch_summary(&self, days: u32) -> Result<UsageSummary> {
        // Try provider APIs first if admin keys are configured.
        // Merge Anthropic + OpenAI summaries when both are present.
        let mut api_models: Vec<ModelUsage> = Vec::new();
        let mut used_api = false;

        if let Some(key) = self.admin_key("anthropic") {
            match provider_api::fetch_anthropic(&key, days).await {
                Ok(s) => {
                    api_models.extend(s.by_model);
                    used_api = true;
                }
                Err(e) => tracing::warn!("Anthropic API failed: {}", e),
            }
        }
        if let Some(key) = self.admin_key("openai") {
            match provider_api::fetch_openai(&key, days).await {
                Ok(s) => {
                    api_models.extend(s.by_model);
                    used_api = true;
                }
                Err(e) => tracing::warn!("OpenAI API failed: {}", e),
            }
        }

        if used_api {
            // Apply pricing config to any rows the provider API didn't price itself
            // (OpenAI returns tokens but not cost on the usage endpoint).
            for m in &mut api_models {
                if m.cost_usd == 0.0 && (m.input_tokens > 0 || m.output_tokens > 0) {
                    m.cost_usd = self.model_cost(
                        &m.model,
                        TokenCounts {
                            input: m.input_tokens,
                            output: m.output_tokens,
                            cache_creation_input: m.cache_creation_input_tokens,
                            cache_read_input: m.cache_read_input_tokens,
                        },
                    );
                }
            }
            let total = api_models.iter().map(|m| m.cost_usd).sum::<f64>();
            for m in &mut api_models {
                m.pct = if total > 0.0 { m.cost_usd / total * 100.0 } else { 0.0 };
            }
            api_models.sort_by(|a, b| {
                b.cost_usd.partial_cmp(&a.cost_usd).unwrap_or(std::cmp::Ordering::Equal)
            });
            let routing_savings_usd = self.compute_savings(days).unwrap_or(0.0);
            return Ok(UsageSummary {
                days,
                by_model: api_models,
                total_cost_usd: total,
                routing_savings_usd,
                source: StatsSource::ProviderApi,
            });
        }

        // Fallback: parse ~/.claude/ session logs
        let raw = scrape_claude_logs(days)?;
        if raw.is_empty() {
            return Ok(UsageSummary {
                days,
                by_model: vec![],
                total_cost_usd: 0.0,
                routing_savings_usd: 0.0,
                source: StatsSource::Unavailable,
            });
        }

        let by_model = self.cost_by_model(raw);
        let total = by_model.iter().map(|m| m.cost_usd).sum::<f64>();
        let routing_savings_usd = self.compute_savings(days).unwrap_or(0.0);

        Ok(UsageSummary {
            days,
            by_model,
            total_cost_usd: total,
            routing_savings_usd,
            source: StatsSource::LocalLogScraper,
        })
    }

    fn admin_key(&self, provider: &str) -> Option<String> {
        let cfg = match provider {
            "anthropic" => self.config.providers.anthropic.as_ref()?,
            "openai" => self.config.providers.openai.as_ref()?,
            _ => return None,
        };
        let k = cfg.admin_key.clone()?;
        if k.is_empty() { None } else { Some(k) }
    }

    /// Returns estimated USD saved by routing in the last `days`.
    /// Opens its own DB connection so the engine doesn't need to hold one.
    fn compute_savings(&self, days: u32) -> Result<f64> {
        let db_path = self.config.daemon.db_path();
        if !db_path.exists() {
            return Ok(0.0);
        }
        let db = Db::open(&db_path)?;
        savings::compute(&db.conn, &self.config, days)
    }

    fn cost_by_model(&self, raw: HashMap<String, TokenCounts>) -> Vec<ModelUsage> {
        let total_cost: f64 = raw
            .iter()
            .map(|(model, counts)| self.model_cost(model, *counts))
            .sum();

        let mut usage: Vec<ModelUsage> = raw
            .into_iter()
            .map(|(model, counts)| {
                let cost = self.model_cost(&model, counts);
                ModelUsage {
                    model,
                    input_tokens: counts.input,
                    output_tokens: counts.output,
                    cache_creation_input_tokens: counts.cache_creation_input,
                    cache_read_input_tokens: counts.cache_read_input,
                    cost_usd: cost,
                    pct: if total_cost > 0.0 { cost / total_cost * 100.0 } else { 0.0 },
                }
            })
            .collect();

        usage.sort_by(|a, b| {
            b.cost_usd.partial_cmp(&a.cost_usd).unwrap_or(std::cmp::Ordering::Equal)
        });
        usage
    }

    fn model_cost(&self, model: &str, c: TokenCounts) -> f64 {
        if let Some(pricing) = self.config.pricing.get(model) {
            let per = |tokens: u64, rate: f64| (tokens as f64 / 1_000_000.0) * rate;
            per(c.input, pricing.input)
                + per(c.output, pricing.output)
                + per(c.cache_creation_input, pricing.cache_creation_or_default())
                + per(c.cache_read_input, pricing.cache_read_or_default())
        } else {
            // Unknown model: conservative mid-tier estimate (no cache pricing)
            let per = |tokens: u64, rate: f64| (tokens as f64 / 1_000_000.0) * rate;
            per(c.input, 3.0) + per(c.output, 15.0) + per(c.cache_read_input, 0.30)
        }
    }
}