#!/usr/bin/env python3
"""GRPO Stage-B terminal dashboard.

Run:
  watch -n 15 -c .venv/bin/python scripts/grpo_watch.py
or:
  .venv/bin/python scripts/grpo_watch.py --loop 15
"""
from __future__ import annotations

import argparse
import json
import math
import os
import shutil
import statistics as st
import subprocess
import sys
import time
from pathlib import Path

RUN_DIR = Path("runs/grpo_stageb")
HISTORY = RUN_DIR / "history.jsonl"
CANARY = RUN_DIR / "canary.jsonl"
LOG = Path("runs/grpo_stageb.log")  # launch log sits beside the run dir
WORKER_LOG = RUN_DIR / "worker.log"

# ANSI
RESET = "\033[0m"
BOLD = "\033[1m"
DIM = "\033[2m"
RED = "\033[31m"
GREEN = "\033[32m"
YELLOW = "\033[33m"
BLUE = "\033[34m"
MAGENTA = "\033[35m"
CYAN = "\033[36m"
WHITE = "\033[37m"
BG_RED = "\033[41m"
BG_GREEN = "\033[42m"
BG_YELLOW = "\033[43m"
BG_BLUE = "\033[44m"
BG_MAGENTA = "\033[45m"
BG_CYAN = "\033[46m"


def fmt(x, nd=4, na="—"):
    if x is None:
        return na
    if isinstance(x, float) and math.isnan(x):
        return "NaN"
    return f"{x:.{nd}f}"


def bar(cur, total, width=28, fill="█", empty="░", color=GREEN):
    if total <= 0:
        return f"{DIM}{empty * width}{RESET}"
    frac = min(max(cur / total, 0.0), 1.0)
    n = int(frac * width)
    return f"{color}{fill * n}{DIM}{empty * (width - n)}{RESET}"


def sparkline(values, width=40):
    if not values:
        return ""
    blocks = "▁▂▃▄▅▆▇█"
    vals = list(values)[-width:]
    lo, hi = min(vals), max(vals)
    rng = (hi - lo) or 1.0
    return "".join(blocks[min(int((v - lo) / rng * (len(blocks) - 1)), len(blocks) - 1)] for v in vals)


def read_jsonl(path: Path):
    if not path.exists():
        return []
    out = []
    with open(path) as f:
        for line in f:
            line = line.strip()
            if line:
                try:
                    out.append(json.loads(line))
                except Exception:
                    continue
    return out


def tail(path: Path, n=6):
    if not path.exists():
        return []
    try:
        with open(path, "rb") as f:
            f.seek(0, os.SEEK_END)
            end = f.tell()
            size = min(end, 8192)
            f.seek(end - size)
            data = f.read().decode(errors="replace")
        return data.splitlines()[-n:]
    except Exception:
        return []


def gpu_lines():
    try:
        out = subprocess.check_output(
            [
                "nvidia-smi",
                "--query-gpu=index,utilization.gpu,memory.used,memory.total,temperature.gpu",
                "--format=csv,noheader,nounits",
            ],
            text=True,
        )
    except Exception:
        return ["GPU: nvidia-smi unavailable"]
    lines = []
    for row in out.strip().splitlines():
        idx, util, used, total, temp = [c.strip() for c in row.split(",")]
        used_g = float(used) / 1024
        total_g = float(total) / 1024
        lines.append(
            f"GPU{idx} {bar(float(util), 100, width=14, color=CYAN)} {util:>3}% "
            f"{used_g:5.1f}/{total_g:.0f}G {temp}°C"
        )
    return lines


def proc_status():
    torchrun = subprocess.run(["pgrep", "-f", "torchrun.*grpo"], capture_output=True, text=True).stdout.strip()
    vllm = subprocess.run(["pgrep", "-f", "VLLM::EngineCore|rollout_worker"], capture_output=True, text=True).stdout.strip()
    return bool(torchrun), bool(vllm)


def render():
    cols = shutil.get_terminal_size((100, 30)).columns
    rows = read_jsonl(HISTORY)
    canary = read_jsonl(CANARY)
    torchrun_up, vllm_up = proc_status()

    opt_steps = [r.get("optimizer_step", 0) for r in rows]
    cur_opt = max(opt_steps) if opt_steps else 0
    micro_steps = [r.get("micro_step", 0) for r in rows]
    cur_micro = max(micro_steps) if micro_steps else 0
    grad_accum = 85  # configs/grpo_stageb.toml
    micro_in_step = cur_micro - cur_opt * grad_accum
    frac_opt = cur_opt + micro_in_step / grad_accum  # fractional step for the bar
    applied = [r for r in rows if r.get("grad_norm") == r.get("grad_norm")]
    gn = [r["grad_norm"] for r in applied]
    acc = [r["reward/accuracy"] for r in rows if r.get("reward/accuracy") is not None]
    ent = [r["entropy"] for r in rows if r.get("entropy") is not None]
    loss = [r["loss"] for r in rows if r.get("loss") is not None]
    lr = [r["lr"] for r in rows if r.get("lr") is not None]
    resp = [r["lengths/response_mean"] for r in rows if r.get("lengths/response_mean") is not None]

    total_steps = 500
    # run start = first timestamped line of the launch log (ctime is NOT
    # creation time on Linux -- it tracks the last write, which would
    # make elapsed/pace collapse to ~0 while the run is healthy).
    run_start = None
    if LOG.exists():
        try:
            with open(LOG) as f:
                for line in f:
                    if " INFO " in line and line[:4].isdigit():
                        run_start = time.mktime(time.strptime(line[:19], "%Y-%m-%d %H:%M:%S"))
                        break
        except Exception:
            pass
    if run_start is None:
        run_start = HISTORY.stat().st_mtime if HISTORY.exists() else time.time()
    elapsed = time.time() - run_start

    # Pace from observed micro-step rate (state file), not run start:
    # run start includes one-time boot (engine load, canary, first rollout),
    # which makes early pace readings garbage.
    state_f = RUN_DIR / ".watch_state.json"
    try:
        wstate = json.loads(state_f.read_text()) if state_f.exists() else {}
    except Exception:
        wstate = {}
    now = time.time()
    ref_micro, ref_ts = wstate.get("ref", (None, None))
    pace_min = wstate.get("pace_min")
    if ref_micro is None:
        wstate["ref"] = (cur_micro, now)
        pace_min = None
    elif now - ref_ts > 180 and cur_micro > ref_micro:
        micros_per_min = (cur_micro - ref_micro) / ((now - ref_ts) / 60)
        # 85 training micros + one rollout generation per step
        pace_min = grad_accum / micros_per_min + 2.5
        wstate["ref"] = (cur_micro, now)
        wstate["pace_min"] = pace_min
    try:
        state_f.write_text(json.dumps(wstate))
    except Exception:
        pass
    eta_min = pace_min * (total_steps - frac_opt) if pace_min else None

    print(f"{BOLD}{CYAN}╔{'═' * (cols - 2)}╗{RESET}")
    title = f" GRPO Stage-B • optimizer step {cur_opt}/{total_steps} "
    print(f"{BOLD}{CYAN}║{title.center(cols - 2)}║{RESET}")
    print(f"{BOLD}{CYAN}╚{'═' * (cols - 2)}╝{RESET}")
    print()

    # Staleness alarm: history must be written every few minutes while alive
    hist_age = (time.time() - HISTORY.stat().st_mtime) if HISTORY.exists() else None
    stale = (hist_age is None) or (hist_age > 900)
    if stale and torchrun_up:
        print(f"  {BG_RED}{BOLD}{WHITE} ⚠ WEDGE? history.jsonl untouched for "
              f"{(hist_age or 0)/60:.0f} min while torchrun is alive — check the log {RESET}")
        print()

    # Status line
    status = f"{GREEN}RUNNING{RESET}" if torchrun_up else f"{BG_RED}{BOLD}{WHITE} STOPPED — NO TRAINING PROCESS {RESET}"
    engine = f"{GREEN}vLLM up{RESET}" if vllm_up else f"{YELLOW}vLLM down{RESET}"
    print(f"  status: {status}  |  engine: {engine}  |  rows: {len(rows)}  |  applied steps: {len(applied)}")
    print(f"  progress: {bar(frac_opt, total_steps, width=cols - 22)} {frac_opt:.1f}/{total_steps}")
    print(f"  step {cur_opt}: micro {micro_in_step}/{grad_accum} {bar(micro_in_step, grad_accum, width=20, color=YELLOW)}")
    if pace_min:
        print(f"  pace: {pace_min:.1f} min/step  |  elapsed: {elapsed/60:.0f} min  |  ETA: {eta_min/60:.1f} h")
    else:
        print(f"  pace: measuring… (needs a few min of micro history)  |  elapsed: {elapsed/60:.0f} min")
    print()

    # Two-column metric panels
    left_w = cols // 2 - 3
    right_w = cols - left_w - 6

    def panel(title, lines, width):
        out = [f"{BOLD}{MAGENTA}┌─ {title} {'─' * max(0, width - len(title) - 4)}┐{RESET}"]
        for ln in lines:
            out.append(f"{MAGENTA}│{RESET} {ln.ljust(width - 2)} {MAGENTA}│{RESET}")
        out.append(f"{MAGENTA}└{'─' * (width - 2)}┘{RESET}")
        return out

    train_lines = [
        f"loss      {fmt(loss[-1] if loss else None)}  {sparkline(loss, 24)}",
        f"grad_norm {fmt(gn[-1] if gn else None)}  med {fmt(st.median(gn) if gn else None)}  max {fmt(max(gn) if gn else None)}",
        f"entropy   {fmt(ent[-1] if ent else None)}  {sparkline(ent, 24)}",
        f"lr        {fmt(lr[-1] if lr else None, 8)}",
        f"micro     {cur_micro}",
    ]
    reward_lines = [
        f"acc       {fmt(acc[-1] if acc else None)}  mean {fmt(st.mean(acc) if acc else None)}",
        f"acc first50 {fmt(st.mean(acc[:50]) if len(acc) >= 50 else None)}  last50 {fmt(st.mean(acc[-50:]) if len(acc) >= 50 else None)}",
        f"resp_len  {fmt(resp[-1] if resp else None, 1)}  {sparkline(resp, 24)}",
        f"canary    {canary[-1].get('canary/pass@1') if canary else '—'} (n={canary[-1].get('canary/n') if canary else '—'})",
        f"canary step {canary[-1].get('optimizer_step') if canary else '—'}",
    ]

    left = panel("Training", train_lines, left_w)
    right = panel("Reward / Canary", reward_lines, right_w)
    for l, r in zip(left, right):
        print(f"  {l}  {r}")
    print()

    # GPU panel
    print(f"  {BOLD}{BLUE}GPUs{RESET}")
    for ln in gpu_lines():
        print(f"    {ln}")
    print()

    # Recent log tail
    print(f"  {BOLD}{BLUE}Recent log (grpo_stageb.log){RESET}")
    for ln in tail(LOG, 5):
        print(f"    {DIM}{ln[: cols - 6]}{RESET}")
    print()
    print(f"  {DIM}updated {time.strftime('%H:%M:%S')}  |  Ctrl-C to exit{RESET}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--loop", type=int, default=0, help="refresh seconds; 0 = render once")
    args = ap.parse_args()
    if args.loop <= 0:
        render()
        return
    try:
        while True:
            # clear screen without shell (avoid os.system injection sink)
            sys.stdout.write("\033[2J\033[H")
            sys.stdout.flush()
            render()
            time.sleep(args.loop)
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()