"""Shared token/cache usage aggregation over Claude Code transcripts.

Single source of truth for the per-model in/out/cache breakdown rendered by
both the status line (statusline-cost.sh, whole-session) and the post-run
workflow usage report hook (workflow-usage-report.py, time-windowed to one run).

Pass window=None for whole-file accumulation (status line). Pass
window=(t0_ms, t1_ms) to keep only assistant turns whose ISO timestamp falls
in [t0_ms, t1_ms] epoch-milliseconds (one workflow run).
"""
import os, json
from datetime import datetime

PRICING = {
    "opus":   (15.00, 75.00, 1.50, 18.75),
    "sonnet": (3.00,  15.00, 0.30,  3.75),
    "haiku":  (0.80,   4.00, 0.08,  1.00),
    "fable":  (30.00, 150.00, 3.00, 37.50),
}
RENDER_ORDER = ["fable", "opus", "sonnet", "haiku"]
NAMES = {"opus": "Opus", "sonnet": "Sonnet", "haiku": "Haiku", "fable": "Fable"}
LABEL_COLORS = {"opus": "\033[00;35m", "sonnet": "\033[00;36m", "haiku": "\033[00;33m", "fable": "\033[00;31m"}
DIM = "\033[02m"
RESET = "\033[00m"


def model_family(model_id):
    m = (model_id or "").lower()
    if "opus" in m: return "opus"
    if "sonnet" in m: return "sonnet"
    if "haiku" in m: return "haiku"
    if "fable" in m: return "fable"
    return "other"


def _iso_to_ms(ts):
    try:
        return datetime.fromisoformat(ts.replace("Z", "+00:00")).timestamp() * 1000.0
    except Exception:
        return None


def parse_files(paths, window=None):
    """Accumulate per-model usage across jsonl transcript files.

    Returns (per_model, first_ts, last_ts, turn_count) where per_model maps a
    model family to {"in","out","cache_read","cache_write"} and first/last_ts
    are ISO strings. window, if given, is (t0_ms, t1_ms): assistant entries
    outside the inclusive window are skipped entirely (not counted).
    """
    per_model = {}
    first_ts = None
    last_ts = None
    turn_count = 0
    for path in paths:
        if not path or not os.path.exists(path):
            continue
        try:
            with open(path) as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        entry = json.loads(line)
                    except Exception:
                        continue
                    if entry.get("type") != "assistant":
                        continue
                    ts = entry.get("timestamp")
                    if window is not None:
                        ms = _iso_to_ms(ts) if ts else None
                        if ms is None or ms < window[0] or ms > window[1]:
                            continue
                    turn_count += 1
                    msg = entry.get("message", {})
                    model_id = msg.get("model", "")
                    usage = msg.get("usage", {})
                    if not model_id or not usage:
                        continue
                    fam = model_family(model_id)
                    if fam not in per_model:
                        per_model[fam] = {"in": 0, "out": 0, "cache_read": 0, "cache_write": 0}
                    per_model[fam]["in"]          += usage.get("input_tokens", 0)
                    per_model[fam]["out"]         += usage.get("output_tokens", 0)
                    per_model[fam]["cache_read"]  += usage.get("cache_read_input_tokens", 0)
                    per_model[fam]["cache_write"] += usage.get("cache_creation_input_tokens", 0)
                    if ts:
                        if last_ts is None or ts > last_ts:
                            last_ts = ts
                        if first_ts is None or ts < first_ts:
                            first_ts = ts
        except Exception:
            pass
    return per_model, first_ts, last_ts, turn_count


def fmt(n):
    if n >= 1_000_000:
        return f"{n/1_000_000:.1f}M"
    if n >= 1000:
        return f"{n//1000}k"
    return str(n)


def render_parts(per_model, color=True):
    """One "Name I:.. O:.. C:↓.. ↑.. hit:..%" string per present model family,
    in RENDER_ORDER. color=True wraps in the status-line ANSI; color=False is
    plain text for hook/systemMessage output."""
    parts = []
    for fam in RENDER_ORDER:
        if fam not in per_model:
            continue
        u = per_model[fam]
        name = NAMES[fam]
        total_in = u["in"] + u["cache_read"]
        hit_pct = f" hit:{u['cache_read']*100//total_in}%" if total_in > 0 else ""
        detail = f"I:{fmt(u['in'])} O:{fmt(u['out'])} C:↓{fmt(u['cache_read'])} ↑{fmt(u['cache_write'])}{hit_pct}"
        if color:
            c = LABEL_COLORS.get(fam, "\033[00;36m")
            parts.append(f"{c}{name}{RESET} {DIM}{detail}{RESET}")
        else:
            parts.append(f"{name} {detail}")
    return parts


def fmt_duration(secs):
    secs = int(secs)
    if secs >= 86400:
        d = secs // 86400; h = (secs % 86400) // 3600; m = (secs % 3600) // 60
        return f"{d}d{h:02d}h{m:02d}m"
    if secs >= 3600:
        return f"{secs//3600}h{(secs%3600)//60}m"
    if secs >= 60:
        return f"{secs//60}m{secs%60}s"
    return f"{secs}s"