claude-code-stats 0.1.0

pub mod source;
pub mod types;

use std::fs;
use std::path::PathBuf;

use anyhow::Result;
use chrono::{DateTime, Utc};

use source::cli_probe::CliProbeSource;
use source::oauth::OauthSource;
use source::web::WebSource;
use source::{UsageSource, fetch_with_failover};
use types::{
    ActivePayload, ApiResponse, CachedResponse, CostData, DataSource, ExtraUsage, WidgetPayload,
    parse_reset_time, to_usage_window,
};

pub const CACHE_MAX_AGE_SECS: i64 = 4 * 60;
pub const CACHE_FILE_NAME: &str = "claude-code-stats-cache.json";

pub fn collect_widget_payload() -> Result<WidgetPayload> {
    let now = Utc::now();
    let (usage, cached_at, data_source, cost_data) = get_usage_with_cache(now)?;

    let data = ActivePayload {
        five_hour: usage
            .five_hour
            .map(|w| to_usage_window(&w, now, Some(300.0))),
        seven_day: usage
            .seven_day
            .map(|w| to_usage_window(&w, now, Some(10080.0))),
        seven_day_sonnet: usage
            .seven_day_sonnet
            .map(|w| to_usage_window(&w, now, Some(10080.0))),
        seven_day_opus: usage
            .seven_day_opus
            .map(|w| to_usage_window(&w, now, Some(10080.0))),
        extra_usage: usage.extra_usage.map(|e| ExtraUsage {
            is_enabled: e.is_enabled,
            monthly_limit: e.monthly_limit,
            used_credits: e.used_credits,
            utilization: e.utilization,
        }),
        updated_at: cached_at.to_rfc3339(),
        source: data_source,
        cost_data,
    };

    Ok(WidgetPayload::Active {
        data: Box::new(data),
    })
}

pub fn collect_widget_payload_json() -> String {
    let payload = match collect_widget_payload() {
        Ok(payload) => payload,
        Err(err) => {
            eprintln!("claude-code-stats error: {err:?}");
            WidgetPayload::Error {
                title: "Claude usage error".to_string(),
                message: err.to_string(),
            }
        }
    };

    match serde_json::to_string(&payload) {
        Ok(json) => json,
        Err(err) => {
            eprintln!("claude-code-stats serialization error: {err:?}");
            "{\"status\":\"error\",\"title\":\"claude-code-stats\",\"message\":\"serialization failure\"}".to_string()
        }
    }
}

fn get_cache_path() -> PathBuf {
    dirs::cache_dir()
        .unwrap_or_else(|| PathBuf::from("/tmp"))
        .join(CACHE_FILE_NAME)
}

fn get_usage_with_cache(
    now: DateTime<Utc>,
) -> Result<(ApiResponse, DateTime<Utc>, DataSource, Option<CostData>)> {
    let cache_path = get_cache_path();

    if let Ok(contents) = fs::read_to_string(&cache_path) {
        if let Ok(cached) = serde_json::from_str::<CachedResponse>(&contents) {
            if let Some(cached_at) = parse_reset_time(&cached.cached_at) {
                let age = now.signed_duration_since(cached_at).num_seconds();
                if age < CACHE_MAX_AGE_SECS {
                    return Ok((cached.response, cached_at, cached.source, cached.cost_data));
                }
            }
        }
    }

    let sources: Vec<&dyn UsageSource> = vec![&OauthSource, &WebSource, &CliProbeSource];
    let (response, source_name) = fetch_with_failover(&sources)?;

    let data_source = match source_name {
        "oauth_api" => DataSource::OauthApi,
        "web_api" => DataSource::WebApi,
        "cli_probe" => DataSource::CliProbe,
        _ => DataSource::OauthApi,
    };

    let cost_data = source::cost_scanner::scan_cost_data();

    let cached = CachedResponse {
        cached_at: now.to_rfc3339(),
        response: response.clone(),
        source: data_source,
        cost_data: cost_data.clone(),
    };

    if let Ok(json) = serde_json::to_string(&cached) {
        let _ = fs::write(&cache_path, json);
    }

    Ok((response, now, data_source, cost_data))
}