#!/usr/bin/env python3
import sys, json, os, glob
from datetime import datetime, timezone

sys.dont_write_bytecode = True
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "lib"))
from usage_agg import parse_files, render_parts, fmt_duration, PRICING

try:
    data = json.load(sys.stdin)
except Exception:
    sys.exit(0)

transcript_path = data.get("transcript_path", "")
session_id = data.get("session_id", "")

paths = [transcript_path]
if transcript_path and session_id:
    subagents_dir = os.path.join(os.path.dirname(transcript_path), session_id, "subagents")
    paths += sorted(glob.glob(os.path.join(subagents_dir, "agent-*.jsonl")))

per_model, first_ts, last_ts, turn_count = parse_files(paths)
transcript_derived = bool(per_model)

# Turn count: derive from the transcript when present; otherwise count
# distinct prompt_ids seen for this session (one per user turn) in a small
# sidecar log, since the hook payload has no turn counter of its own.
if not turn_count:
    prompt_id = data.get("prompt_id", "")
    turn_count = 0
    if session_id and prompt_id:
        turns_log = f"/home/user/.claude/session-env/{session_id}.turns.log"
        try:
            seen = []
            if os.path.exists(turns_log):
                with open(turns_log) as f:
                    seen = [l.strip() for l in f if l.strip()]
            if not seen or seen[-1] != prompt_id:
                with open(turns_log, "a") as f:
                    f.write(prompt_id + "\n")
                seen.append(prompt_id)
            turn_count = len(set(seen))
        except Exception:
            turn_count = 0

# Model breakdown: fall back to the current session's cumulative context
# window figures (single active model) when no transcript was parseable.
if not per_model:
    model_id = (data.get("model") or {}).get("id", "")
    cw = data.get("context_window") or {}
    cur = cw.get("current_usage") or {}
    total_in = cw.get("total_input_tokens", 0)
    total_out = cw.get("total_output_tokens", 0)
    if model_id and (total_in or total_out):
        from usage_agg import model_family
        fam = model_family(model_id)
        per_model[fam] = {
            "in": total_in,
            "out": total_out,
            "cache_read": cur.get("cache_read_input_tokens", 0),
            "cache_write": cur.get("cache_creation_input_tokens", 0),
        }

# Cost: prefer the whole-session transcript sum (persists across resume);
# else the hook's cumulative total_cost_usd; else price the context-window
# snapshot as a last resort.
hook_cost = (data.get("cost") or {}).get("total_cost_usd")
if transcript_derived or hook_cost is None:
    total_cost = 0.0
    for fam, u in per_model.items():
        p = PRICING.get(fam, (3.00, 15.00, 0.30, 3.75))
        total_cost += (
            u["in"]          / 1_000_000 * p[0] +
            u["out"]         / 1_000_000 * p[1] +
            u["cache_read"]  / 1_000_000 * p[2] +
            u["cache_write"] / 1_000_000 * p[3]
        )
else:
    total_cost = float(hook_cost)

DIM   = "\033[02m"
RESET = "\033[00m"
GREEN  = "\033[00;32m"
YELLOW = "\033[00;33m"
CYAN   = "\033[00;36m"
BOLD_RED = "\033[01;31m"
STALE_SECS = 3600

cost_str = f"${total_cost:.2f}"
turns_str = f"{turn_count}t"

parts = render_parts(per_model, color=True)

last_ts_str = ""
last_ts_color = YELLOW
last_ts_dt = None
if last_ts:
    try:
        last_ts_dt = datetime.fromisoformat(last_ts.replace("Z", "+00:00"))
        last_ts_str = last_ts_dt.astimezone().strftime("%H:%M:%S %d/%m")
    except Exception:
        pass
elif "total_duration_ms" in (data.get("cost") or {}):
    # No transcript timestamps available; show "now" as the last-turn marker.
    last_ts_dt = datetime.now(timezone.utc)
    last_ts_str = last_ts_dt.astimezone().strftime("%H:%M:%S %d/%m")

if last_ts_dt is not None:
    elapsed = (datetime.now(timezone.utc) - last_ts_dt).total_seconds()
    if elapsed > STALE_SECS:
        last_ts_color = BOLD_RED

duration_str = ""
if first_ts and last_ts and first_ts != last_ts:
    try:
        t0 = datetime.fromisoformat(first_ts.replace("Z", "+00:00"))
        t1 = datetime.fromisoformat(last_ts.replace("Z", "+00:00"))
        duration_str = fmt_duration((t1 - t0).total_seconds())
    except Exception:
        pass
elif not (first_ts and last_ts):
    total_dur_ms = (data.get("cost") or {}).get("total_duration_ms")
    if total_dur_ms:
        duration_str = fmt_duration(total_dur_ms / 1000)

line1 = f"{GREEN}{cost_str}{RESET}  {CYAN}{turns_str}{RESET}"
if last_ts_str:
    line1 += f"  {last_ts_color}{last_ts_str}{RESET}"
if duration_str:
    line1 += f"  {DIM}{duration_str}{RESET}"
if line1:
    print(line1, end="")

if parts:
    print("\n" + "  ".join(parts), end="")