"""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, hashlib, tempfile, time
from datetime import datetime

CACHE_DIR = os.path.join(
    os.environ.get("XDG_CACHE_HOME") or os.path.join(os.path.expanduser("~"), ".cache"),
    "claude-usage-agg",
)
CACHE_SCHEMA = 1
CACHE_TTL_S = 14 * 86400

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 _new_acc():
    return {"per_model": {}, "first_ts": None, "last_ts": None, "turns": 0}


def _merge(dst, src):
    for fam, u in src["per_model"].items():
        d = dst["per_model"].setdefault(fam, {"in": 0, "out": 0, "cache_read": 0, "cache_write": 0})
        for k in d:
            d[k] += u.get(k, 0)
    dst["turns"] += src["turns"]
    for key, better in (("first_ts", min), ("last_ts", max)):
        if src[key] is not None:
            dst[key] = src[key] if dst[key] is None else better(dst[key], src[key])


def _scan(f, start, window, acc):
    """Fold every complete JSONL line from byte `start` into acc; return the
    byte offset just past the last complete line (a partially written trailing
    line is left for the next call)."""
    f.seek(start)
    offset = start
    for raw in f:
        if not raw.endswith(b"\n"):
            break
        offset += len(raw)
        line = raw.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
        acc["turns"] += 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)
        u = acc["per_model"].setdefault(fam, {"in": 0, "out": 0, "cache_read": 0, "cache_write": 0})
        u["in"]          += usage.get("input_tokens", 0)
        u["out"]         += usage.get("output_tokens", 0)
        u["cache_read"]  += usage.get("cache_read_input_tokens", 0)
        u["cache_write"] += usage.get("cache_creation_input_tokens", 0)
        if ts:
            if acc["last_ts"] is None or ts > acc["last_ts"]:
                acc["last_ts"] = ts
            if acc["first_ts"] is None or ts < acc["first_ts"]:
                acc["first_ts"] = ts
    return offset


def _cache_path(path, window):
    key = json.dumps([CACHE_SCHEMA, os.path.abspath(path), window], sort_keys=True)
    return os.path.join(CACHE_DIR, hashlib.sha256(key.encode()).hexdigest()[:32] + ".json")


def _cache_store(cache_file, state, prune):
    os.makedirs(CACHE_DIR, exist_ok=True)
    fd, tmp = tempfile.mkstemp(dir=CACHE_DIR, suffix=".tmp")
    with os.fdopen(fd, "w") as f:
        json.dump(state, f)
    os.replace(tmp, cache_file)
    if prune:
        cutoff = time.time() - CACHE_TTL_S
        for name in os.listdir(CACHE_DIR):
            stale = os.path.join(CACHE_DIR, name)
            try:
                if os.path.getmtime(stale) < cutoff:
                    os.unlink(stale)
            except OSError:
                pass


def _file_usage(path, window):
    """Per-file accumulator, resumed from the append-only cache when the file is
    the same inode grown in place; any other change forces a full reparse."""
    st = os.stat(path)
    cache_file = _cache_path(path, window)
    acc, start, cached = _new_acc(), 0, None
    try:
        with open(cache_file) as f:
            cached = json.load(f)
    except Exception:
        pass
    if (
        cached
        and cached.get("dev") == st.st_dev
        and cached.get("ino") == st.st_ino
        and 0 <= cached.get("offset", -1) <= st.st_size
    ):
        acc, start = cached["acc"], cached["offset"]
    with open(path, "rb") as f:
        offset = _scan(f, start, window, acc)
    if offset != start or cached is None:
        _cache_store(
            cache_file,
            {"dev": st.st_dev, "ino": st.st_ino, "offset": offset, "acc": acc},
            prune=cached is None,
        )
    return acc


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).
    """
    total = _new_acc()
    for path in paths:
        if not path or not os.path.exists(path):
            continue
        try:
            _merge(total, _file_usage(path, window))
        except Exception:
            pass
    return total["per_model"], total["first_ts"], total["last_ts"], total["turns"]


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"