codex_usage_profiler/core.py

from __future__ import annotations

import json
import sqlite3
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Callable, Iterable
from urllib.parse import quote
from zoneinfo import ZoneInfo


@dataclass(slots=True)
class TokenBreakdown:
    input_tokens: int = 0
    cached_input_tokens: int = 0
    output_tokens: int = 0
    reasoning_output_tokens: int = 0
    total_tokens: int = 0

    @classmethod
    def from_usage(cls, usage: dict[str, Any]) -> "TokenBreakdown":
        return cls(
            input_tokens=_int_value(usage.get("input_tokens")),
            cached_input_tokens=_int_value(usage.get("cached_input_tokens")),
            output_tokens=_int_value(usage.get("output_tokens")),
            reasoning_output_tokens=_int_value(usage.get("reasoning_output_tokens")),
            total_tokens=_int_value(usage.get("total_tokens")),
        )

    def add(self, other: "TokenBreakdown | None") -> None:
        if other is None:
            return
        self.input_tokens += other.input_tokens
        self.cached_input_tokens += other.cached_input_tokens
        self.output_tokens += other.output_tokens
        self.reasoning_output_tokens += other.reasoning_output_tokens
        self.total_tokens += other.total_tokens


@dataclass(slots=True)
class RolloutStats:
    path: Path | None
    found: bool
    token_breakdown: TokenBreakdown | None = None
    model: str | None = None
    cwd: str | None = None
    turn_count: int = 0
    malformed_lines: int = 0
    warnings: list[str] = field(default_factory=list)


@dataclass(slots=True)
class SessionUsage:
    thread_id: str
    created_at: datetime
    updated_at: datetime
    cwd: str
    model: str
    model_provider: str
    role: str
    agent_role: str
    agent_nickname: str
    kind: str
    parent_thread_id: str | None
    title: str
    tokens: int
    rollout_path: Path | None
    rollout_found: bool
    rollout_breakdown: TokenBreakdown | None
    turn_count: int

    @property
    def week_label(self) -> str:
        year, week, _ = self.created_at.isocalendar()
        return f"{year}-W{week:02d}"

    @property
    def day_label(self) -> str:
        return self.created_at.date().isoformat()


@dataclass(slots=True)
class Bucket:
    label: str
    sessions: int = 0
    tokens: int = 0
    rollout: TokenBreakdown = field(default_factory=TokenBreakdown)

    def add(self, session: SessionUsage) -> None:
        self.sessions += 1
        self.tokens += session.tokens
        self.rollout.add(session.rollout_breakdown)


@dataclass(slots=True)
class UsageReport:
    generated_at: datetime
    timezone_name: str
    state_path: Path
    sessions_dir: Path
    sessions: list[SessionUsage]
    daily: list[Bucket]
    weekly: list[Bucket]
    by_cwd: list[Bucket]
    by_model: list[Bucket]
    by_role: list[Bucket]
    by_kind: list[Bucket]
    top_sessions: list[SessionUsage]
    rollout_totals: TokenBreakdown
    warnings: list[str]
    days_filter: int | None = None

    @property
    def total_tokens(self) -> int:
        return sum(session.tokens for session in self.sessions)

    @property
    def parsed_rollout_count(self) -> int:
        return sum(1 for session in self.sessions if session.rollout_breakdown is not None)

    @property
    def missing_rollout_count(self) -> int:
        return sum(1 for session in self.sessions if not session.rollout_found)

    @property
    def date_range_label(self) -> str:
        if not self.sessions:
            return "n/a"
        first = min(session.created_at for session in self.sessions)
        last = max(session.created_at for session in self.sessions)
        return f"{first.date().isoformat()} to {last.date().isoformat()}"


def build_report(
    *,
    state_path: Path,
    sessions_dir: Path,
    tz_name: str,
    days: int | None = None,
    top: int = 20,
    now: datetime | None = None,
) -> UsageReport:
    tz = ZoneInfo(tz_name)
    generated_at = (now or datetime.now(tz)).astimezone(tz)
    rows, child_to_parent = _load_threads(state_path)
    sessions: list[SessionUsage] = []
    warnings: list[str] = []

    if days is not None and days <= 0:
        raise ValueError("--days must be greater than zero")
    cutoff_date = None
    if days is not None:
        cutoff_date = generated_at.date() - timedelta(days=days - 1)

    for row in rows:
        created_at = _local_datetime(row.get("created_at_ms"), row.get("created_at"), tz)
        if cutoff_date is not None and created_at.date() < cutoff_date:
            continue

        updated_at = _local_datetime(row.get("updated_at_ms"), row.get("updated_at"), tz)
        thread_id = str(row["id"])
        rollout_path = _resolve_rollout_path(row.get("rollout_path"), sessions_dir, thread_id)
        rollout = parse_rollout(rollout_path)
        warnings.extend(rollout.warnings)

        parent_thread_id = child_to_parent.get(thread_id)
        kind = "subagent" if parent_thread_id else "main"
        agent_role = str(row.get("agent_role") or "")
        role = agent_role or ("main" if kind == "main" else "unknown")
        model_provider = str(row.get("model_provider") or "unknown")
        model = str(row.get("model") or rollout.model or "unknown")
        cwd = str(row.get("cwd") or rollout.cwd or "unknown")

        sessions.append(
            SessionUsage(
                thread_id=thread_id,
                created_at=created_at,
                updated_at=updated_at,
                cwd=cwd,
                model=model,
                model_provider=model_provider,
                role=role,
                agent_role=agent_role,
                agent_nickname=str(row.get("agent_nickname") or ""),
                kind=kind,
                parent_thread_id=parent_thread_id,
                title=str(row.get("title") or ""),
                tokens=_int_value(row.get("tokens_used")),
                rollout_path=rollout.path,
                rollout_found=rollout.found,
                rollout_breakdown=rollout.token_breakdown,
                turn_count=rollout.turn_count,
            )
        )

    sessions.sort(key=lambda session: (session.created_at, session.thread_id))
    rollout_totals = TokenBreakdown()
    for session in sessions:
        rollout_totals.add(session.rollout_breakdown)

    return UsageReport(
        generated_at=generated_at,
        timezone_name=tz_name,
        state_path=Path(state_path).expanduser(),
        sessions_dir=Path(sessions_dir).expanduser(),
        sessions=sessions,
        daily=_group_sessions(sessions, lambda session: session.day_label, sort_by_label=True),
        weekly=_group_sessions(sessions, lambda session: session.week_label, sort_by_label=True),
        by_cwd=_group_sessions(sessions, lambda session: _compact_path(session.cwd)),
        by_model=_group_sessions(sessions, lambda session: session.model or "unknown"),
        by_role=_group_sessions(sessions, lambda session: session.role or "unknown"),
        by_kind=_group_sessions(sessions, lambda session: session.kind),
        top_sessions=sorted(sessions, key=lambda session: session.tokens, reverse=True)[:top],
        rollout_totals=rollout_totals,
        warnings=warnings,
        days_filter=days,
    )


def parse_rollout(path: Path | None) -> RolloutStats:
    if path is None:
        return RolloutStats(path=None, found=False, warnings=["No rollout path for a session"])

    expanded = path.expanduser()
    if not expanded.exists():
        return RolloutStats(
            path=expanded,
            found=False,
            warnings=[f"Missing rollout file: {expanded}"],
        )

    stats = RolloutStats(path=expanded, found=True)
    try:
        with expanded.open("r", encoding="utf-8") as handle:
            for line_number, line in enumerate(handle, start=1):
                if not line.strip():
                    continue
                try:
                    event = json.loads(line)
                except json.JSONDecodeError:
                    stats.malformed_lines += 1
                    continue
                _read_rollout_event(stats, event)
    except OSError as exc:
        stats.found = False
        stats.warnings.append(f"Could not read rollout file {expanded}: {exc}")
        return stats

    if stats.malformed_lines:
        stats.warnings.append(
            f"Malformed JSONL lines in {expanded}: {stats.malformed_lines}"
        )
    return stats


def _read_rollout_event(stats: RolloutStats, event: dict[str, Any]) -> None:
    event_type = event.get("type")
    payload = event.get("payload")
    if not isinstance(payload, dict):
        payload = {}

    if event_type == "session_meta":
        stats.cwd = stats.cwd or _string_or_none(payload.get("cwd"))
        return

    if event_type == "turn_context":
        stats.model = stats.model or _string_or_none(payload.get("model"))
        stats.cwd = stats.cwd or _string_or_none(payload.get("cwd"))
        return

    if event_type != "event_msg":
        return

    payload_type = payload.get("type")
    if payload_type == "task_started":
        stats.turn_count += 1
        return
    if payload_type != "token_count":
        return

    info = payload.get("info")
    if not isinstance(info, dict):
        return
    usage = info.get("total_token_usage")
    if isinstance(usage, dict):
        stats.token_breakdown = TokenBreakdown.from_usage(usage)


def _load_threads(state_path: Path) -> tuple[list[dict[str, Any]], dict[str, str]]:
    expanded = Path(state_path).expanduser()
    if not expanded.exists():
        raise FileNotFoundError(f"State database not found: {expanded}")

    uri = f"file:{quote(str(expanded), safe='/')}?mode=ro"
    with sqlite3.connect(uri, uri=True) as conn:
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA query_only = ON")
        rows = [
            dict(row)
            for row in conn.execute(
                """
                SELECT
                    id,
                    rollout_path,
                    created_at,
                    updated_at,
                    created_at_ms,
                    updated_at_ms,
                    model_provider,
                    cwd,
                    title,
                    tokens_used,
                    agent_nickname,
                    agent_role,
                    model
                FROM threads
                ORDER BY created_at_ms ASC, id ASC
                """
            )
        ]
        edges = [
            dict(edge)
            for edge in conn.execute(
                """
                SELECT parent_thread_id, child_thread_id
                FROM thread_spawn_edges
                """
            )
        ]
    return rows, {str(edge["child_thread_id"]): str(edge["parent_thread_id"]) for edge in edges}


def _resolve_rollout_path(
    rollout_path: Any, sessions_dir: Path, thread_id: str
) -> Path | None:
    sessions_root = Path(sessions_dir).expanduser()
    if rollout_path:
        candidate = Path(str(rollout_path)).expanduser()
        if not candidate.is_absolute():
            candidate = sessions_root / candidate
        if candidate.exists():
            return candidate

    if sessions_root.exists():
        matches = sorted(sessions_root.rglob(f"*{thread_id}.jsonl"))
        if matches:
            return matches[0]

    if rollout_path:
        return candidate
    return None


def _group_sessions(
    sessions: Iterable[SessionUsage],
    key: Callable[[SessionUsage], str],
    *,
    sort_by_label: bool = False,
) -> list[Bucket]:
    buckets: dict[str, Bucket] = {}
    for session in sessions:
        label = key(session)
        bucket = buckets.setdefault(label, Bucket(label=label))
        bucket.add(session)

    values = list(buckets.values())
    if sort_by_label:
        values.sort(key=lambda bucket: bucket.label)
    else:
        values.sort(key=lambda bucket: bucket.tokens, reverse=True)
    return values


def _local_datetime(ms_value: Any, seconds_value: Any, tz: ZoneInfo) -> datetime:
    if ms_value is not None:
        timestamp = _int_value(ms_value) / 1000
    else:
        timestamp = _int_value(seconds_value)
    return datetime.fromtimestamp(timestamp, timezone.utc).astimezone(tz)


def _compact_path(value: str) -> str:
    if not value:
        return "unknown"
    try:
        path = Path(value).expanduser()
        home = Path.home()
        if path.is_absolute():
            relative = path.relative_to(home)
            return "~" if str(relative) == "." else "~/" + str(relative)
        return str(path)
    except ValueError:
        return value


def _int_value(value: Any) -> int:
    if value is None:
        return 0
    try:
        return int(value)
    except (TypeError, ValueError):
        return 0


def _string_or_none(value: Any) -> str | None:
    if value is None:
        return None
    text = str(value)
    return text if text else None