codex_usage.py

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
import json
import sqlite3
import time


WINDOW_SECONDS = 5 * 60 * 60
_DEFAULT_STATE_DB = object()

PLUS_MESSAGE_RANGES = {
    "gpt-5.5": (15, 80),
    "gpt-5.4": (20, 100),
    "gpt-5.4 mini": (60, 350),
    "gpt-5.4-mini": (60, 350),
}


@dataclass(frozen=True)
class UsageEvent:
    timestamp: int
    model: str


@dataclass(frozen=True)
class UsageSnapshot:
    message_count: int
    model_counts: dict[str, int] = field(default_factory=dict)
    tokens_used: int = 0
    reset_seconds: int | None = None
    lower_limit: int | None = None
    upper_limit: int | None = None

    @property
    def risk_percent(self) -> int | None:
        if not self.lower_limit:
            return None
        return min(999, round((self.message_count / self.lower_limit) * 100))

    @property
    def progress_percent(self) -> int | None:
        if not self.upper_limit:
            return None
        return min(100, round((self.message_count / self.upper_limit) * 100))


@dataclass(frozen=True)
class SnapshotText:
    status: str
    used: str
    recovery: str
    progress_percent: int


def default_sessions_dir() -> Path:
    return Path.home() / ".codex" / "sessions"


def default_state_db() -> Path:
    return Path.home() / ".codex" / "state_5.sqlite"


def read_local_usage(
    sessions_dir: Path | None = None,
    state_db: Path | None | object = _DEFAULT_STATE_DB,
) -> tuple[list[UsageEvent], int]:
    sessions_dir = sessions_dir or default_sessions_dir()
    state_db = default_state_db() if state_db is _DEFAULT_STATE_DB else state_db
    return _read_session_events(sessions_dir), _read_recent_tokens(state_db)


def compute_snapshot(
    events: tuple[list[UsageEvent], int] | list[UsageEvent],
    now: int | None = None,
    window_seconds: int = WINDOW_SECONDS,
) -> UsageSnapshot:
    now = int(time.time()) if now is None else now
    if isinstance(events, tuple):
        usage_events, tokens_used = events
    else:
        usage_events, tokens_used = events, 0

    cutoff = now - window_seconds
    recent = [event for event in usage_events if cutoff <= event.timestamp <= now]
    model_counts: dict[str, int] = {}
    for event in recent:
        model_counts[event.model] = model_counts.get(event.model, 0) + 1

    lower, upper = _weighted_message_range(model_counts)
    reset_seconds = None
    if recent:
        oldest = min(event.timestamp for event in recent)
        reset_seconds = max(0, oldest + window_seconds - now)

    return UsageSnapshot(
        message_count=len(recent),
        model_counts=model_counts,
        tokens_used=tokens_used,
        reset_seconds=reset_seconds,
        lower_limit=lower,
        upper_limit=upper,
    )


def _read_session_events(sessions_dir: Path) -> list[UsageEvent]:
    if not sessions_dir.exists():
        return []

    events: list[UsageEvent] = []
    for path in sessions_dir.rglob("*.jsonl"):
        current_model = "unknown"
        try:
            with path.open("r", encoding="utf-8") as handle:
                for line in handle:
                    try:
                        row = json.loads(line)
                    except json.JSONDecodeError:
                        continue
                    payload = row.get("payload") or {}
                    current_model = _model_from_payload(payload) or current_model
                    if row.get("type") != "event_msg":
                        continue
                    if payload.get("type") != "user_message":
                        continue
                    timestamp = _parse_timestamp(row.get("timestamp"))
                    if timestamp is not None:
                        events.append(UsageEvent(timestamp, current_model))
        except OSError:
            continue
    return events


def _read_recent_tokens(state_db: Path | None) -> int:
    if not state_db or not state_db.exists():
        return 0
    cutoff = int(time.time()) - WINDOW_SECONDS
    try:
        con = sqlite3.connect(f"file:{state_db}?mode=ro", uri=True, timeout=1)
        try:
            row = con.execute(
                "select coalesce(sum(tokens_used), 0) from threads where updated_at >= ?",
                (cutoff,),
            ).fetchone()
            return int(row[0] or 0)
        finally:
            con.close()
    except sqlite3.Error:
        return 0


def _model_from_payload(payload: dict) -> str | None:
    if isinstance(payload.get("model"), str):
        return payload["model"]
    if isinstance(payload.get("session_id"), str) and isinstance(payload.get("model"), str):
        return payload["model"]
    return None


def _parse_timestamp(value: object) -> int | None:
    if not isinstance(value, str):
        return None
    normalized = value.replace("Z", "+00:00")
    try:
        return int(datetime.fromisoformat(normalized).timestamp())
    except ValueError:
        return None


def _weighted_message_range(model_counts: dict[str, int]) -> tuple[int | None, int | None]:
    if not model_counts:
        return None, None
    lower = 0
    upper = 0
    for model, count in model_counts.items():
        model_lower, model_upper = PLUS_MESSAGE_RANGES.get(model, (15, 80))
        lower += model_lower * count
        upper += model_upper * count
    total = sum(model_counts.values())
    return round(lower / total), round(upper / total)


def format_duration(seconds: int | None) -> str:
    if seconds is None:
        return "--:--"
    minutes = max(0, seconds) // 60
    return f"{minutes // 60}:{minutes % 60:02d}"


def format_recovery(seconds: int | None) -> str:
    if seconds is None:
        return "\u6682\u65f6\u6ca1\u6709\u672c\u673a\u8bb0\u5f55"
    minutes = max(0, seconds) // 60
    if minutes < 60:
        return f"\u7ea6 {minutes}\u5206 \u540e\u5f00\u59cb\u6062\u590d"
    return f"\u7ea6 {minutes // 60}\u5c0f\u65f6{minutes % 60:02d}\u5206 \u540e\u5f00\u59cb\u6062\u590d"


def format_snapshot_text(snapshot: UsageSnapshot) -> SnapshotText:
    progress = snapshot.progress_percent or 0
    capped = min(100, progress)
    if progress >= 90:
        status = "\u5feb\u649e\u4e86"
    elif progress >= 70:
        status = "\u5feb\u6ee1\u4e86"
    elif progress >= 40:
        status = "\u8fd8\u884c"
    else:
        status = "\u5f88\u5b89\u5168"
    return SnapshotText(
        status=status,
        used=f"\u5df2\u7528 {snapshot.message_count} \u6b21",
        recovery=format_recovery(snapshot.reset_seconds),
        progress_percent=capped,
    )