cortex-rs-stats 0.2.0

use anyhow::Result;
use chrono::{Duration, Utc};
use serde_json::Value;
use std::collections::HashMap;
use std::path::PathBuf;

#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct TokenCounts {
    pub input: u64,
    pub output: u64,
    pub cache_creation_input: u64,
    pub cache_read_input: u64,
}

/// Scan ~/.claude/ JSONL session files for usage data within `days`.
/// Returns map of model → TokenCounts (input/output/cache_creation/cache_read).
pub fn scrape_claude_logs(days: u32) -> Result<HashMap<String, TokenCounts>> {
    let claude_dir = claude_dir();
    if !claude_dir.exists() {
        return Ok(HashMap::new());
    }

    let cutoff = Utc::now() - Duration::days(days as i64);
    let cutoff_ts = cutoff.timestamp();
    let mut totals: HashMap<String, TokenCounts> = HashMap::new();

    scan_dir(&claude_dir, cutoff_ts, &mut totals)?;

    Ok(totals)
}

fn scan_dir(
    dir: &std::path::Path,
    cutoff_ts: i64,
    totals: &mut HashMap<String, TokenCounts>,
) -> Result<()> {
    let entries = match std::fs::read_dir(dir) {
        Ok(e) => e,
        Err(_) => return Ok(()),
    };

    for entry in entries.flatten() {
        let path = entry.path();
        if path.is_dir() {
            scan_dir(&path, cutoff_ts, totals)?;
        } else if path.extension().map(|e| e == "jsonl").unwrap_or(false) {
            // File-level short circuit: if the whole file hasn't been touched
            // since the cutoff, no line in it can be in-window.
            if let Ok(meta) = entry.metadata() {
                if let Ok(modified) = meta.modified() {
                    if let Ok(elapsed) = modified.duration_since(std::time::UNIX_EPOCH) {
                        if (elapsed.as_secs() as i64) < cutoff_ts {
                            continue;
                        }
                    }
                }
            }
            parse_jsonl(&path, cutoff_ts, totals)?;
        }
    }
    Ok(())
}

fn parse_jsonl(
    path: &std::path::Path,
    cutoff_ts: i64,
    totals: &mut HashMap<String, TokenCounts>,
) -> Result<()> {
    let content = match std::fs::read_to_string(path) {
        Ok(c) => c,
        Err(_) => return Ok(()),
    };

    for line in content.lines() {
        let line = line.trim();
        if line.is_empty() {
            continue;
        }
        let Ok(v): Result<Value, _> = serde_json::from_str(line) else {
            continue;
        };

        // Real Claude logs use two timestamp formats — int milliseconds in
        // history.jsonl, and ISO 8601 strings in session JSONLs. Parse both.
        // If neither parses, skip the entry conservatively (don't count untimestamped rows).
        let Some(ts) = parse_timestamp(v.get("timestamp")) else {
            continue;
        };
        if ts < cutoff_ts {
            continue;
        }

        extract_usage(&v, totals);
    }

    Ok(())
}

/// Accepts: int seconds, int milliseconds, ISO 8601 string. Returns unix seconds.
fn parse_timestamp(v: Option<&Value>) -> Option<i64> {
    let v = v?;
    if let Some(i) = v.as_i64() {
        // Disambiguate ms vs s by magnitude. Anything > 1e12 is ms (year 33658 in s).
        return Some(if i > 1_000_000_000_000 { i / 1000 } else { i });
    }
    if let Some(s) = v.as_str() {
        return chrono::DateTime::parse_from_rfc3339(s)
            .ok()
            .map(|dt| dt.timestamp());
    }
    None
}

fn extract_usage(v: &Value, totals: &mut HashMap<String, TokenCounts>) {
    if let (Some(model), Some(usage)) = (
        v.get("model").and_then(|m| m.as_str()),
        v.get("usage"),
    ) {
        let input = usage.get("input_tokens").and_then(|t| t.as_u64()).unwrap_or(0);
        let output = usage.get("output_tokens").and_then(|t| t.as_u64()).unwrap_or(0);
        let cache_creation = usage
            .get("cache_creation_input_tokens")
            .and_then(|t| t.as_u64())
            .unwrap_or(0);
        let cache_read = usage
            .get("cache_read_input_tokens")
            .and_then(|t| t.as_u64())
            .unwrap_or(0);

        if input + output + cache_creation + cache_read > 0 {
            let entry = totals.entry(model.to_string()).or_default();
            entry.input += input;
            entry.output += output;
            entry.cache_creation_input += cache_creation;
            entry.cache_read_input += cache_read;
            return;
        }
    }

    for key in &["message", "result", "response"] {
        if let Some(nested) = v.get(key) {
            extract_usage(nested, totals);
        }
    }
}

fn claude_dir() -> PathBuf {
    let home = std::env::var("HOME").unwrap_or_default();
    PathBuf::from(home).join(".claude")
}

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

    #[test]
    fn parses_usage_from_jsonl_line() {
        let mut totals = HashMap::new();
        let line = r#"{"timestamp": 9999999999, "model": "claude-sonnet-4-6", "usage": {"input_tokens": 1000, "output_tokens": 500}}"#;
        let v: Value = serde_json::from_str(line).unwrap();
        extract_usage(&v, &mut totals);
        let counts = totals.get("claude-sonnet-4-6").copied().unwrap();
        assert_eq!(counts.input, 1000);
        assert_eq!(counts.output, 500);
        assert_eq!(counts.cache_creation_input, 0);
        assert_eq!(counts.cache_read_input, 0);
    }

    #[test]
    fn parses_cache_tokens() {
        let mut totals = HashMap::new();
        let line = r#"{"timestamp": 9999999999, "model": "claude-opus-4-7", "usage": {"input_tokens": 6, "output_tokens": 170, "cache_creation_input_tokens": 44102, "cache_read_input_tokens": 0}}"#;
        let v: Value = serde_json::from_str(line).unwrap();
        extract_usage(&v, &mut totals);
        let c = totals.get("claude-opus-4-7").copied().unwrap();
        assert_eq!(c.input, 6);
        assert_eq!(c.output, 170);
        assert_eq!(c.cache_creation_input, 44102);
    }

    #[test]
    fn skips_old_entries() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("session.jsonl");
        let mut f = std::fs::File::create(&path).unwrap();
        // timestamp 0 = very old
        writeln!(f, r#"{{"timestamp": 0, "model": "claude-opus-4-7", "usage": {{"input_tokens": 9999, "output_tokens": 9999}}}}"#).unwrap();

        let mut totals = HashMap::new();
        parse_jsonl(&path, chrono::Utc::now().timestamp() - 100, &mut totals).unwrap();
        assert!(totals.is_empty(), "old entry should be skipped");
    }

    #[test]
    fn parses_iso_timestamp() {
        let ts = parse_timestamp(Some(&serde_json::json!("2026-05-13T14:25:07.145Z")));
        assert!(ts.is_some());
        // 2026-05-13T14:25:07Z = 1778682307 (computed via python)
        assert_eq!(ts.unwrap(), 1778682307);
    }

    #[test]
    fn parses_int_milliseconds() {
        let ts = parse_timestamp(Some(&serde_json::json!(1778386184520_i64)));
        assert_eq!(ts.unwrap(), 1778386184);
    }

    #[test]
    fn parses_int_seconds() {
        let ts = parse_timestamp(Some(&serde_json::json!(1778386184_i64)));
        assert_eq!(ts.unwrap(), 1778386184);
    }

    #[test]
    fn untimestamped_entries_excluded() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("session.jsonl");
        let mut f = std::fs::File::create(&path).unwrap();
        // No timestamp field — must NOT be counted
        writeln!(f, r#"{{"model": "claude-opus-4-7", "usage": {{"input_tokens": 9999, "output_tokens": 9999}}}}"#).unwrap();

        let mut totals = HashMap::new();
        parse_jsonl(&path, 0, &mut totals).unwrap();
        assert!(totals.is_empty(), "untimestamped row leaked through cutoff filter");
    }

    #[test]
    fn iso_timestamp_in_window_is_counted() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("session.jsonl");
        let mut f = std::fs::File::create(&path).unwrap();
        let recent = chrono::Utc::now() - chrono::Duration::hours(1);
        let iso = recent.to_rfc3339();
        writeln!(
            f,
            r#"{{"timestamp": "{}", "model": "claude-sonnet-4-6", "usage": {{"input_tokens": 100, "output_tokens": 50}}}}"#,
            iso
        ).unwrap();

        let mut totals = HashMap::new();
        let cutoff = (chrono::Utc::now() - chrono::Duration::days(7)).timestamp();
        parse_jsonl(&path, cutoff, &mut totals).unwrap();
        let c = totals.get("claude-sonnet-4-6").copied().unwrap();
        assert_eq!(c.input, 100);
        assert_eq!(c.output, 50);
    }
}