#!/usr/bin/env python3
import argparse
import json
import re
import statistics
import subprocess
import time
from collections import Counter, defaultdict
from datetime import datetime, timezone

import requests

PREVIEW = {"generate_preview_frames", "gpu_task:preview"}
BAKE = {"texture_model_for_export", "gpu_task:bake"}
RECORD = re.compile(r"(GPU_CALL|LABS_EVENT) (\{.*\})")
ASCTIME = re.compile(r"(\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2})")


def hf_token():
    query = "protocol=https\nhost=huggingface.co\n\n"
    out = subprocess.run(["git", "credential-osxkeychain", "get"], input=query, capture_output=True, text=True).stdout
    return next(line.split("=", 1)[1] for line in out.splitlines() if line.startswith("password="))


def fetch_run_log(space, seconds):
    url = f"https://huggingface.co/api/spaces/broyang/{space}/logs/run"
    lines, deadline = [], time.time() + seconds
    try:
        with requests.get(url, headers={"Authorization": f"Bearer {hf_token()}"}, stream=True, timeout=(10, 5)) as resp:
            resp.raise_for_status()
            for raw in resp.iter_lines(decode_unicode=True):
                if raw:
                    lines.append(raw)
                if time.time() > deadline:
                    break
    except requests.exceptions.ConnectionError:
        pass
    return lines


def unwrap(raw):
    if not raw.startswith("data:"):
        return raw, None
    try:
        event = json.loads(raw[5:].strip())
        return event.get("data", ""), event.get("timestamp")
    except ValueError:
        return raw[5:], None


def parse(lines):
    calls, events, seen, stamps = [], [], set(), []
    for raw in lines:
        text, stamp = unwrap(raw)
        stamped = stamp or (ASCTIME.search(text).group(1) if ASCTIME.search(text) else None)
        if stamped:
            stamps.append(stamped)
        match = RECORD.search(text)
        if not match or match.group(0) in seen:
            continue
        seen.add(match.group(0))
        try:
            record = json.loads(match.group(2))
        except ValueError:
            continue
        (calls if match.group(1) == "GPU_CALL" else events).append(record)
    return sorted(calls, key=lambda r: r.get("t", 0)), sorted(events, key=lambda r: r.get("t", 0)), stamps


def pct(part, whole):
    return f"{part}/{whole} ({100 * part / whole:.0f}%)" if whole else "0/0 (n/a)"


def spread(values):
    if not values:
        return "n/a"
    values = sorted(values)
    pick = lambda q: values[min(len(values) - 1, int(q * len(values)))]
    return f"n={len(values)} min={values[0]:.1f} p25={pick(.25):.1f} p50={pick(.5):.1f} p75={pick(.75):.1f} max={values[-1]:.1f}"


def report_reuse(calls):
    print("== GPU calls ==")
    by_fn = defaultdict(list)
    for call in calls:
        by_fn[call.get("fn")].append(call)
    for fn, rows in sorted(by_fn.items()):
        reused = [r for r in rows if r.get("calls_in_worker", 1) > 1]
        fresh = [r for r in rows if r.get("calls_in_worker", 1) <= 1]
        print(f"{fn}: calls={len(rows)} ok={sum(bool(r.get('ok')) for r in rows)} reused={pct(len(reused), len(rows))}")
        print(f"  since_prev_call_s (reused): {spread([r['since_prev_call_s'] for r in reused if r.get('since_prev_call_s') is not None])}")
        print(f"  in_gpu_s fresh:  {spread([r['in_gpu_s'] for r in fresh if 'in_gpu_s' in r])}")
        print(f"  in_gpu_s reused: {spread([r['in_gpu_s'] for r in reused if 'in_gpu_s' in r])}")
        print(f"  texture_size: {dict(Counter(r.get('texture_size') for r in rows))}")
    pairs = [(prev, cur) for prev, cur in zip(calls, calls[1:]) if prev.get("fn") in PREVIEW and cur.get("fn") in PREVIEW]
    cross = [(p, c) for p, c in pairs if p.get("session") != c.get("session")]
    hit = lambda group: sum(c.get("calls_in_worker", 1) > 1 for _, c in group)
    print(f"preview after preview: reused={pct(hit(pairs), len(pairs))}; different session: reused={pct(hit(cross), len(cross))}")
    bakes = [(p, c) for p, c in zip(calls, calls[1:])
             if p.get("fn") in PREVIEW and c.get("fn") in BAKE and p.get("session") == c.get("session")]
    print(f"bake right after its own preview: reused={pct(hit(bakes), len(bakes))}")
    print(f"  gap preview->next preview start (s): {spread([c['t'] - p['t'] - p.get('in_gpu_s', 0) for p, c in pairs if 't' in c and 't' in p])}")


def report_click_through(events):
    print("== Click-through (LABS_EVENT) ==")
    by_session = defaultdict(list)
    for event in events:
        if event.get("session"):
            by_session[event["session"]].append(event)
    previews = sum(e.get("event") == "preview_done" for e in events)
    print(f"preview_done={previews} (sessions={len(by_session)}, untagged={sum(not e.get('session') for e in events)})")
    for kind in ("glb_done", "stl_done"):
        followed, delays = 0, []
        for rows in by_session.values():
            for index, event in enumerate(rows):
                if event.get("event") != "preview_done":
                    continue
                for later in rows[index + 1:]:
                    if later.get("event") == "preview_done":
                        break
                    if later.get("event") == kind:
                        followed += 1
                        delays.append(later["t"] - later.get("took_s", 0) - event["t"])
                        break
        print(f"{kind.split('_')[0].upper()}: previews followed by a click={pct(followed, previews)}"
              f" median preview->click={f'{statistics.median(delays):.1f}s' if delays else 'n/a'}  {spread(delays)}")
        print(f"  {kind} ok={sum(e.get('event') == kind and bool(e.get('ok')) for e in events)}"
              f" failed={sum(e.get('event') == kind and not e.get('ok') for e in events)}")


def report_window(lines, stamps, calls, events):
    times = [r["t"] for r in calls + events if "t" in r]
    fmt = lambda t: datetime.fromtimestamp(t, timezone.utc).strftime("%Y-%m-%d %H:%M:%SZ")
    print("== Log window ==")
    print(f"lines read={len(lines)}; first log line={stamps[0] if stamps else 'n/a'}; last={stamps[-1] if stamps else 'n/a'}")
    if times:
        print(f"records: {len(calls)} GPU_CALL, {len(events)} LABS_EVENT, {fmt(min(times))} .. {fmt(max(times))}"
              f" ({(max(times) - min(times)) / 3600:.1f} h)")
    print("Note: the HF run log only covers the current container (since its last start/restart).")


def main():
    parser = argparse.ArgumentParser(description="ZeroGPU worker reuse + GLB/STL click-through from 3D AI logs")
    parser.add_argument("--space", default="3dai", help="broyang/<space> (default 3dai = production)")
    parser.add_argument("--file", help="saved log file instead of fetching the live run log")
    parser.add_argument("--seconds", type=float, default=20, help="how long to read the SSE log stream")
    args = parser.parse_args()
    if args.file:
        with open(args.file, encoding="utf-8", errors="replace") as handle:
            lines = handle.read().splitlines()
    else:
        lines = fetch_run_log(args.space, args.seconds)
    calls, events, stamps = parse(lines)
    report_window(lines, stamps, calls, events)
    report_reuse(calls)
    report_click_through(events)


if __name__ == "__main__":
    main()