#!/usr/bin/env python3
"""
Parse a Cursor usage-events CSV into per-pool, per-model token and
API-equivalent cost breakdown.

Pools (Cursor Pro):
  · Auto + Composer  — c25/auto before the auto→API cutover (Pro sub)
  · API credits      — Claude/GPT/etc anytime + c25/auto after cutover
  · Free bonus       — Kind=Free rows
  · No charge        — errored/aborted rows

Usage:
    python3 cursor_usage.py [csv]
    python3 cursor_usage.py --auto-cutoff 2026-06-13T00:00:00Z [csv]
"""

from __future__ import annotations

import argparse
import csv
import datetime as dt
import sys
from collections import defaultdict
from typing import Optional

# (input, cache_write, cache_read, output) — $ per 1M tokens.
# Longer prefixes first so "composer-2.5-fast" isn't caught by "composer-2.5".
PRICING: dict[str, tuple[float, float, float, float]] = {
    "composer-2.5-fast":                 (3.00,   3.00,  0.50,  15.00),
    "composer-2.5":                      (0.50,   0.50,  0.20,   2.50),
    "auto":                              (1.25,   1.25,  0.25,   6.00),
    "claude-fable-5-thinking-high":      (10.00, 12.50,  1.00,  50.00),
    "claude-fable-5":                    (10.00, 12.50,  1.00,  50.00),
    "claude-4.6-sonnet-medium-thinking": (3.00,   3.75,  0.30,  15.00),
    "claude-4.6-sonnet":                 (3.00,   3.75,  0.30,  15.00),
    "claude-opus-4.8":                   (5.00,   6.25,  0.50,  25.00),
    "claude-4.6-opus":                   (5.00,   6.25,  0.50,  25.00),
    "gemini-3.1-pro":                    (2.00,   2.00,  0.20,  12.00),
    "gemini-3.5-flash":                  (1.50,   1.50,  0.15,   9.00),
    "gpt-5.3-codex":                     (1.75,   1.75,  0.175, 14.00),
    "gpt-5.5":                           (5.00,   5.00,  0.50,  30.00),
    "grok-build-0.1":                    (1.00,   1.00,  0.20,   2.00),
}

AUTO_CUTOFF_FALLBACK = "2026-06-13T00:00:00.000Z"
UNPRICED_TAG = "(unpriced)"
NO_CHARGE_KINDS = {"Errored, No Charge", "Aborted, Not Charged"}
AUTO_COMPOSER_PREFIXES = ("auto", "composer-2.5")

POOLS = ("auto_composer", "api_credit", "free_bonus", "no_charge")
POOL_LABELS = {
    "auto_composer": "Auto + Composer pool  (Pro subscription)",
    "api_credit":    "API credit pool",
    "free_bonus":    "Free bonus  (Kind = Free)",
    "no_charge":     "Errored / aborted  (no charge)",
}

Stat = dict[str, int | float]


def load_events(path: str) -> list[dict]:
    with open(path, newline="", encoding="utf-8") as f:
        rows = list(csv.DictReader(f))
    rows.sort(key=lambda r: r["Date"])
    return rows


def _int_field(event: dict, key: str) -> int:
    return int((event.get(key) or "0").strip() or 0)


def parse_tokens(event: dict) -> dict[str, int]:
    return {
        "in_wcache": _int_field(event, "Input (w/ Cache Write)"),
        "in":        _int_field(event, "Input (w/o Cache Write)"),
        "cache":     _int_field(event, "Cache Read"),
        "out":       _int_field(event, "Output Tokens"),
    }


def price_for(model: str) -> tuple[str, Optional[tuple[float, float, float, float]]]:
    m = (model or "").lower().strip()
    if m in PRICING:
        return m, PRICING[m]
    for key, rate in PRICING.items():
        if m.startswith(key):
            return key, rate
    return model or UNPRICED_TAG, None


def api_equiv_cost(
    tokens: dict[str, int], rate: Optional[tuple[float, float, float, float]]
) -> float:
    if rate is None:
        return 0.0
    ri, rw, rc, ro = rate
    return (
        tokens["in_wcache"] / 1e6 * rw
        + tokens["in"]        / 1e6 * ri
        + tokens["cache"]     / 1e6 * rc
        + tokens["out"]       / 1e6 * ro
    )


def is_auto_composer(model: str) -> bool:
    m = (model or "").lower()
    return any(m.startswith(p) for p in AUTO_COMPOSER_PREFIXES)


def detect_auto_cutoff(
    events: list[dict],
    min_neighbors: int = 30,
    min_after: int = 5,
) -> tuple[str, str]:
    """Find when c25/auto stopped drawing from the Auto pool."""
    daily: dict[str, int] = defaultdict(int)
    for e in events:
        if e["Kind"] not in ("Included", "Free"):
            continue
        if not is_auto_composer(e["Model"]):
            continue
        daily[e["Date"][:10]] += 1

    if not daily:
        return AUTO_CUTOFF_FALLBACK, "no c25/auto events — using fallback date"

    days = sorted(daily)
    candidates: list[tuple[str, int, int, int]] = []

    for i in range(len(days) - 1):
        d_prev, d_next = days[i], days[i + 1]
        gap = (
            dt.datetime.strptime(d_next, "%Y-%m-%d")
            - dt.datetime.strptime(d_prev, "%Y-%m-%d")
        ).days
        if gap > 1:
            gap_start = (
                dt.datetime.strptime(d_prev, "%Y-%m-%d") + dt.timedelta(days=1)
            ).strftime("%Y-%m-%d")
            candidates.append((gap_start, daily[d_prev], daily[d_next], gap - 1))

    qualified = [
        c for c in candidates if c[1] >= min_neighbors and c[2] >= min_after
    ]
    if not qualified:
        return AUTO_CUTOFF_FALLBACK, "no qualifying gap — using fallback date"

    # Latest qualifying gap: auto pool usually drains late in the period.
    gap_start, prev_n, next_n, gap_len = max(qualified, key=lambda c: (c[0], c[1] + c[2]))
    cutover = f"{gap_start}T00:00:00.000Z"
    reason = (
        f"{gap_start}  ({gap_len}-day gap after {prev_n}-event day, "
        f"resumed with {next_n} events)"
    )
    return cutover, reason


def classify(events: list[dict], cutover: str) -> list[dict]:
    rows: list[dict] = []
    for e in events:
        tokens = parse_tokens(e)
        model_key, rate = price_for(e["Model"])
        cost = api_equiv_cost(tokens, rate)
        kind = e["Kind"]

        if kind == "Free":
            pool = "free_bonus"
        elif kind in NO_CHARGE_KINDS:
            pool, cost = "no_charge", 0.0
        elif is_auto_composer(e["Model"]) and e["Date"] < cutover:
            pool = "auto_composer"
        else:
            pool = "api_credit"

        total = tokens["in_wcache"] + tokens["in"] + tokens["cache"] + tokens["out"]
        rows.append({
            "pool": pool,
            "model": model_key,
            "events": 1,
            **tokens,
            "total": total,
            "cost": cost,
        })
    return rows


def aggregate(rows: list[dict]) -> dict[tuple[str, str], Stat]:
    out: dict[tuple[str, str], Stat] = defaultdict(
        lambda: {"events": 0, "in_wcache": 0, "in": 0, "cache": 0, "out": 0, "total": 0, "cost": 0.0}
    )
    for r in rows:
        k = (r["pool"], r["model"])
        for field in ("events", "in_wcache", "in", "cache", "out", "total"):
            out[k][field] += r[field]  # type: ignore[operator]
        out[k]["cost"] += r["cost"]  # type: ignore[operator]
    return out


def fmt_tokens(n: int | float) -> str:
    n = int(n)
    if n >= 1_000_000_000:
        return f"{n / 1e9:.2f}B"
    if n >= 1_000_000:
        return f"{n / 1e6:.2f}M"
    if n >= 1_000:
        return f"{n / 1e3:.1f}K"
    return str(n)


def fmt_money(n: float) -> str:
    if n >= 100:
        return f"${n:,.2f}"
    return f"${n:.2f}"


def empty_stat() -> Stat:
    return {"events": 0, "in_wcache": 0, "in": 0, "cache": 0, "out": 0, "total": 0, "cost": 0.0}


def add_stat(a: Stat, b: Stat) -> Stat:
    return {k: a[k] + b[k] for k in a}  # type: ignore[operator]


def render(
    by_pool_model: dict[tuple[str, str], Stat],
    events: list[dict],
    cutover: str,
    cutover_note: str,
) -> None:
    period_start = events[0]["Date"][:10]
    period_end = events[-1]["Date"][:10]
    model_w = 36
    width = 96

    print()
    print("═" * width)
    print("  Cursor Usage Report")
    print("═" * width)
    print(f"  Period     {period_start}  →  {period_end}   ({len(events):,} events)")
    print(f"  Cutover    {cutover[:10]}   {cutover_note}")
    print("═" * width)

    grand = empty_stat()
    header = (
        f"  {'Model':<{model_w}} {'Events':>7} {'Input':>9} {'Cache':>9} "
        f"{'Output':>9} {'Total':>9} {'API-equiv':>10}"
    )

    for pool in POOLS:
        models = sorted(m for (p, m) in by_pool_model if p == pool)
        print()
        print(f"  {POOL_LABELS[pool]}")
        print("  " + "─" * (width - 2))

        if not models:
            print("  (none)")
            continue

        print(header)
        print("  " + "─" * (width - 2))

        sub = empty_stat()
        for model in models:
            d = by_pool_model[(pool, model)]
            print(
                f"  {model:<{model_w}} {d['events']:>7,} "
                f"{fmt_tokens(d['in'] + d['in_wcache']):>9} "
                f"{fmt_tokens(d['cache']):>9} "
                f"{fmt_tokens(d['out']):>9} "
                f"{fmt_tokens(d['total']):>9} "
                f"{fmt_money(d['cost']):>10}"
            )
            sub = add_stat(sub, d)

        print("  " + "─" * (width - 2))
        print(
            f"  {'Subtotal':<{model_w}} {sub['events']:>7,} "
            f"{fmt_tokens(sub['in'] + sub['in_wcache']):>9} "
            f"{fmt_tokens(sub['cache']):>9} "
            f"{fmt_tokens(sub['out']):>9} "
            f"{fmt_tokens(sub['total']):>9} "
            f"{fmt_money(sub['cost']):>10}"
        )
        grand = add_stat(grand, sub)

    print()
    print("═" * width)
    print(
        f"  {'Grand total':<{model_w}} {grand['events']:>7,} "
        f"{fmt_tokens(grand['in'] + grand['in_wcache']):>9} "
        f"{fmt_tokens(grand['cache']):>9} "
        f"{fmt_tokens(grand['out']):>9} "
        f"{fmt_tokens(grand['total']):>9} "
        f"{fmt_money(grand['cost']):>10}"
    )
    print("═" * width)
    print()
    print("  API-equiv = cache-inclusive cost at published list prices.")
    print()


def main() -> int:
    p = argparse.ArgumentParser(description="Cursor usage-events CSV breakdown.")
    p.add_argument("csv", nargs="?", default="usage-events-2026-06-22.csv")
    p.add_argument("--auto-cutoff", metavar="ISO", help="Override auto→API cutover timestamp.")
    p.add_argument("--no-auto-detect", action="store_true", help="Skip cutover detection.")
    args = p.parse_args()

    events = load_events(args.csv)

    if args.auto_cutoff:
        cutover, note = args.auto_cutoff, "user override"
    elif args.no_auto_detect:
        cutover, note = AUTO_CUTOFF_FALLBACK, "fallback (detection disabled)"
    else:
        cutover, note = detect_auto_cutoff(events)

    render(aggregate(classify(events, cutover)), events, cutover, note)
    return 0


if __name__ == "__main__":
    sys.exit(main())