tokr 0.1.0

Persistent token-usage ledger for AI coding agents. Captures on write, queries forever.
use anyhow::Result;
use chrono::{DateTime, Duration, Local, Utc};
use clap::Args;
use rusqlite::types::Value;

use crate::db::Db;
use crate::ui;

const BLOCK_HOURS: i64 = 5;

#[derive(Args, Debug)]
pub struct BlocksArgs {
    #[arg(long, value_enum)]
    pub source: Option<SourceFilter>,
    #[arg(long)]
    pub model: Option<String>,
    #[arg(long)]
    pub project: Option<String>,
    #[arg(long, value_name = "WINDOW")]
    pub last: Option<String>,
    #[arg(long)]
    pub utc: bool,
    #[arg(long, default_value_t = 20)]
    pub limit: usize,
    #[arg(long)]
    pub json: bool,
    #[arg(long)]
    pub include_synthetic: bool,
}

#[derive(Clone, Copy, Debug, clap::ValueEnum)]
pub enum SourceFilter {
    ClaudeCode,
    Codex,
}

#[derive(Debug, serde::Serialize)]
pub struct Block {
    pub start: DateTime<Utc>,
    pub end_planned: DateTime<Utc>,
    pub last_event: DateTime<Utc>,
    pub events: u64,
    pub tokens: u64,
    pub cost: f64,
    pub models: Vec<String>,
    pub is_active: bool,
}

pub fn run(args: BlocksArgs) -> Result<()> {
    let db = Db::open()?;

    let mut clauses: Vec<String> = Vec::new();
    let mut binds: Vec<Value> = Vec::new();

    if !args.include_synthetic {
        clauses.push("model != '<synthetic>'".into());
    }
    if let Some(s) = args.source {
        let v = match s {
            SourceFilter::ClaudeCode => "claude_code",
            SourceFilter::Codex => "codex",
        };
        clauses.push(format!("source = ?{}", binds.len() + 1));
        binds.push(Value::Text(v.into()));
    }
    if let Some(m) = &args.model {
        clauses.push(format!("model LIKE ?{}", binds.len() + 1));
        binds.push(Value::Text(format!("%{m}%")));
    }
    if let Some(p) = &args.project {
        clauses.push(format!("project_path LIKE ?{}", binds.len() + 1));
        binds.push(Value::Text(format!("%{p}%")));
    }
    if let Some(w) = &args.last {
        if let Some(t) = parse_window(w) {
            clauses.push(format!("timestamp >= ?{}", binds.len() + 1));
            binds.push(Value::Text(t));
        }
    }

    let where_sql = if clauses.is_empty() {
        String::new()
    } else {
        format!("WHERE {}", clauses.join(" AND "))
    };

    let sql = format!(
        "SELECT timestamp, model,
                input_tokens + output_tokens + cache_creation_5m
                  + cache_creation_1h + cache_read_tokens AS tokens,
                cost_usd
         FROM usage_events {where_sql}
         ORDER BY timestamp ASC"
    );

    let mut stmt = db.conn.prepare(&sql)?;
    let rows = stmt
        .query_map(rusqlite::params_from_iter(binds.iter()), |r| {
            Ok((
                r.get::<_, String>(0)?,
                r.get::<_, String>(1)?,
                r.get::<_, i64>(2)? as u64,
                r.get::<_, f64>(3)?,
            ))
        })?
        .collect::<Result<Vec<_>, _>>()?;

    let blocks = compute_blocks(&rows);
    let now = Utc::now();
    let mut blocks = blocks
        .into_iter()
        .map(|mut b| {
            b.is_active = now >= b.start && now < b.end_planned;
            b
        })
        .collect::<Vec<_>>();

    blocks.sort_by_key(|b| std::cmp::Reverse(b.start));
    if args.limit > 0 && blocks.len() > args.limit {
        blocks.truncate(args.limit);
    }

    if args.json {
        println!("{}", serde_json::to_string_pretty(&blocks)?);
    } else {
        print_blocks(&blocks, args.utc);
    }
    Ok(())
}

fn compute_blocks(events: &[(String, String, u64, f64)]) -> Vec<Block> {
    let mut out: Vec<Block> = Vec::new();
    for (ts_str, model, tokens, cost) in events {
        let Ok(ts) = DateTime::parse_from_rfc3339(ts_str) else {
            continue;
        };
        let ts = ts.with_timezone(&Utc);

        let need_new = match out.last() {
            None => true,
            Some(last) => ts >= last.end_planned,
        };

        if need_new {
            out.push(Block {
                start: ts,
                end_planned: ts + Duration::hours(BLOCK_HOURS),
                last_event: ts,
                events: 1,
                tokens: *tokens,
                cost: *cost,
                models: vec![model.clone()],
                is_active: false,
            });
        } else {
            let last = out.last_mut().unwrap();
            last.events += 1;
            last.tokens += tokens;
            last.cost += cost;
            last.last_event = ts;
            if !last.models.iter().any(|m| m == model) {
                last.models.push(model.clone());
            }
        }
    }
    out
}

fn parse_window(s: &str) -> Option<String> {
    let s = s.trim();
    let (n, unit) = if let Some(x) = s.strip_suffix('d') {
        (x, 'd')
    } else if let Some(x) = s.strip_suffix('h') {
        (x, 'h')
    } else if let Some(x) = s.strip_suffix('w') {
        (x, 'w')
    } else {
        (s, 'd')
    };
    let n: i64 = n.parse().ok()?;
    let dur = match unit {
        'h' => Duration::hours(n),
        'w' => Duration::weeks(n),
        _ => Duration::days(n),
    };
    Some((Utc::now() - dur).to_rfc3339())
}

fn print_blocks(blocks: &[Block], utc: bool) {
    if blocks.is_empty() {
        println!(
            "{}  Run {} first or relax the filters.",
            ui::yellow("No blocks found."),
            ui::bold_cyan("tokr sync"),
        );
        return;
    }

    if let Some(active) = blocks.iter().find(|b| b.is_active) {
        render_active_hero(active, utc);
        println!();
    }

    let completed: Vec<&Block> = blocks.iter().filter(|b| !b.is_active).collect();
    if !completed.is_empty() {
        let tz_label = if utc { "UTC" } else { "local" };
        println!(
            "  {} {}   {}",
            ui::bold_cyan("Recent blocks"),
            ui::dim(&format!(
                "· {} block{}",
                completed.len(),
                if completed.len() == 1 { "" } else { "s" }
            )),
            ui::dim(&format!("· tz {tz_label}")),
        );
        println!();

        let mut t = ui::Table::new(
            vec!["Start", "End", "Events", "Tokens", "Cost", "Models"],
            vec![
                ui::Align::Left,
                ui::Align::Left,
                ui::Align::Right,
                ui::Align::Right,
                ui::Align::Right,
                ui::Align::Left,
            ],
        );
        for b in &completed {
            t.push(vec![
                ui::cyan(&fmt_dt(b.start, utc)),
                ui::dim(&fmt_dt(b.end_planned, utc)),
                ui::fmt_int(b.events),
                ui::bold_white(&ui::fmt_compact(b.tokens)),
                ui::green(&ui::fmt_cost(b.cost)),
                ui::magenta(&ui::truncate(
                    &b.models
                        .iter()
                        .map(|m| ui::short_model(m))
                        .collect::<Vec<_>>()
                        .join(", "),
                    44,
                )),
            ]);
        }

        let total_events: u64 = completed.iter().map(|b| b.events).sum();
        let total_tokens: u64 = completed.iter().map(|b| b.tokens).sum();
        let total_cost: f64 = completed.iter().map(|b| b.cost).sum();
        t.with_totals(vec![
            String::from("TOTAL"),
            format!(
                "{} block{}",
                completed.len(),
                if completed.len() == 1 { "" } else { "s" }
            ),
            ui::fmt_int(total_events),
            ui::fmt_compact(total_tokens),
            ui::fmt_cost(total_cost),
            String::new(),
        ]);

        for line in t.render().lines() {
            println!("  {line}");
        }
    }
}

fn render_active_hero(b: &Block, utc: bool) {
    let inner = 76usize;
    let now = Utc::now();
    let elapsed = (now - b.start).num_seconds().max(0) as f64;
    let window = (b.end_planned - b.start).num_seconds() as f64;
    let frac = (elapsed / window).clamp(0.0, 1.0);
    let remaining = b.end_planned - now;
    let rem_mins = remaining.num_minutes().max(0);
    let rem_h = rem_mins / 60;
    let rem_m = rem_mins % 60;

    println!("{}", ui::box_top("Active 5h billing block", inner));
    println!("{}", ui::box_blank(inner));

    let col_w = 22;
    let labels_1 = format!(
        "  {}  {}  {}",
        ui::pad_right(&ui::dim("Started"), col_w),
        ui::pad_right(&ui::dim("Window ends"), col_w),
        ui::pad_right(&ui::dim("Remaining"), col_w),
    );
    let values_1 = format!(
        "  {}  {}  {}",
        ui::pad_right(&ui::cyan(&fmt_dt(b.start, utc)), col_w),
        ui::pad_right(&ui::dim(&fmt_dt(b.end_planned, utc)), col_w),
        ui::pad_right(&ui::bold_white(&format!("{rem_h}h {rem_m:02}m")), col_w),
    );
    let labels_2 = format!(
        "  {}  {}  {}",
        ui::pad_right(&ui::dim("Events"), col_w),
        ui::pad_right(&ui::dim("Tokens"), col_w),
        ui::pad_right(&ui::dim("Cost so far"), col_w),
    );
    let values_2 = format!(
        "  {}  {}  {}",
        ui::pad_right(&ui::bold_white(&ui::fmt_int(b.events)), col_w),
        ui::pad_right(&ui::bold_white(&ui::fmt_compact(b.tokens)), col_w),
        ui::pad_right(&ui::bold_green(&ui::fmt_cost(b.cost)), col_w),
    );

    println!("{}", ui::box_row(&labels_1, inner));
    println!("{}", ui::box_row(&values_1, inner));
    println!("{}", ui::box_blank(inner));
    println!("{}", ui::box_row(&labels_2, inner));
    println!("{}", ui::box_row(&values_2, inner));
    println!("{}", ui::box_blank(inner));

    let bar_w = 50usize;
    let bar = ui::bar_threshold(frac, bar_w);
    let pct_label = format!("{:>5.1}%", frac * 100.0);
    let line = format!(
        "  {}  {}  {}",
        ui::pad_right(&ui::dim("Window"), 8),
        bar,
        ui::bold_white(&pct_label),
    );
    println!("{}", ui::box_row(&line, inner));

    if !b.models.is_empty() {
        let models_line = format!(
            "  {}  {}",
            ui::pad_right(&ui::dim("Models"), 8),
            ui::magenta(&ui::truncate(
                &b.models
                    .iter()
                    .map(|m| ui::short_model(m))
                    .collect::<Vec<_>>()
                    .join(", "),
                60
            )),
        );
        println!("{}", ui::box_row(&models_line, inner));
    }
    println!("{}", ui::box_blank(inner));
    println!("{}", ui::box_bottom(inner));
}

fn fmt_dt(t: DateTime<Utc>, utc: bool) -> String {
    if utc {
        t.format("%Y-%m-%d %H:%M").to_string()
    } else {
        t.with_timezone(&Local).format("%Y-%m-%d %H:%M").to_string()
    }
}