#!/usr/bin/env python3
"""
Gradio control/monitoring UI for the TinyChess research Space.

This is NOT the training implementation. Everything runs from the CLI; this UI
only launches, monitors and inspects background jobs, so a browser disconnect
never kills an experiment.
"""
from __future__ import annotations

import json
import os
import signal
import subprocess
import sys
import time
from pathlib import Path

# --- CPU-ONLY SPACE ---------------------------------------------------------
# This Space previously imported `spaces` and exposed an @spaces.GPU function,
# which made HF classify it as a ZeroGPU Space (requested_hardware=zero-a10g).
# That hardware now requires a PRO subscription, so the Space wedged in
# RUNTIME_ERROR ("403: Subscribe to PRO to use ZeroGPU") and no jobs could run.
# All real research is CPU-only by design (the GPU had a ~5 min cap and was
# never used for training), so the dependency is removed entirely.
class _NoGPU:
    @staticmethod
    def GPU(*a, **k):
        def deco(fn):
            return fn
        return deco(a[0]) if a and callable(a[0]) else deco
spaces = _NoGPU()
_HAS_SPACES = False

import gradio as gr

ROOT = Path(__file__).parent.resolve()
RUNS = ROOT / "runs"
RUNS.mkdir(exist_ok=True)
JOBS = RUNS / "jobs.json"
PY = sys.executable


# ----------------------------------------------------------------- job state
def load_jobs() -> dict:
    if JOBS.exists():
        try:
            return json.loads(JOBS.read_text())
        except Exception:
            return {}
    return {}


def save_jobs(j):
    JOBS.write_text(json.dumps(j, indent=2))


def job_alive(pid) -> bool:
    try:
        os.kill(int(pid), 0)
        return True
    except Exception:
        return False


def _pid_is_ours(pid: int, cmd: list[str]) -> bool:
    """True only if `pid` is alive AND its cmdline matches the recorded job.

    After a Space rebuild the PID space restarts, so a recorded pid can be
    re-used by an unrelated process. Checking liveness alone made launch()
    refuse valid jobs and made job_table() report dead jobs as RUNNING.
    """
    if not job_alive(pid):
        return False
    try:
        with open(f"/proc/{pid}/cmdline", "rb") as fh:
            actual = fh.read().decode(errors="ignore").split("\x00")
    except Exception:
        return False
    want = os.path.basename(cmd[1]) if len(cmd) > 1 else ""
    return any(want and want in a for a in actual)


def launch(name: str, cmd: list[str]) -> str:
    jobs = load_jobs()
    prev = jobs.get(name)
    if prev and _pid_is_ours(prev.get("pid", -1), prev.get("cmd", [])):
        return f"Job '{name}' is already running (pid {prev['pid']})."
    logp = RUNS / f"{name}.log"
    with open(logp, "a") as lf:
        lf.write(f"\n=== launch {time.strftime('%F %T')} ===\n{' '.join(cmd)}\n")
        p = subprocess.Popen(cmd, stdout=lf, stderr=subprocess.STDOUT, cwd=ROOT,
                             start_new_session=True)
    jobs[name] = {"pid": p.pid, "cmd": cmd, "started": time.time(),
                  "log": str(logp)}
    save_jobs(jobs)
    return f"Launched '{name}' (pid {p.pid})\n$ {' '.join(cmd)}"


def stop_job(name: str) -> str:
    jobs = load_jobs()
    if name not in jobs:
        return f"No job '{name}'."
    pid = jobs[name]["pid"]
    try:
        os.killpg(os.getpgid(pid), signal.SIGTERM)
    except Exception as e:
        return f"Could not stop {name}: {e}"
    return f"Sent SIGTERM to '{name}' (pid {pid})."


def job_table() -> str:
    jobs = load_jobs()
    if not jobs:
        return "No jobs yet."
    rows = ["| job | pid | state | runtime |", "|---|---|---|---|"]
    for n, j in sorted(jobs.items()):
        alive = _pid_is_ours(j.get("pid", -1), j.get("cmd", []))
        rt = time.time() - j["started"]
        rows.append(f"| {n} | {j['pid']} | {'RUNNING' if alive else 'stopped'} "
                    f"| {rt/60:.1f} min |")
    return "\n".join(rows)


def tail_log(name: str, n: int = 60) -> str:
    p = RUNS / f"{name}.log"
    if not p.exists():
        return f"(no log for '{name}')"
    lines = p.read_text(errors="replace").splitlines()
    return "\n".join(lines[-n:])


# ----------------------------------------------------------------- metrics
def list_runs():
    return sorted([d.name for d in RUNS.iterdir() if d.is_dir()])


def read_metrics(run: str):
    p = RUNS / run / "metrics.jsonl"
    if not p.exists():
        return [], "(no metrics.jsonl)"
    rows = []
    for line in p.read_text(errors="replace").splitlines():
        try:
            rows.append(json.loads(line))
        except Exception:
            pass
    iters = [r for r in rows if r.get("event") == "iter"]
    return iters, f"{len(iters)} iterations logged"


def metrics_plot(run: str):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    iters, msg = read_metrics(run)
    if not iters:
        return None, msg
    fig, ax = plt.subplots(2, 2, figsize=(10, 6))
    it = [r["iter"] for r in iters]
    ax[0, 0].plot(it, [r.get("sp_decisive_rate", 0) for r in iters]); ax[0, 0].set_title("decisive rate")
    ax[0, 1].plot(it, [r.get("sp_mean_entropy", 0) for r in iters]); ax[0, 1].set_title("policy entropy")
    ax[1, 0].plot(it, [r.get("tr_value_mae", 0) for r in iters]); ax[1, 0].set_title("value MAE")
    ax[1, 1].plot(it, [r.get("sp_mean_plies", 0) for r in iters]); ax[1, 1].set_title("mean plies")
    for a in ax.ravel():
        a.grid(alpha=.3); a.set_xlabel("iteration")
    fig.tight_layout()
    return fig, msg


def summary_table(run: str) -> str:
    iters, msg = read_metrics(run)
    if not iters:
        return msg
    keys = ["iter", "opponent", "sp_decisive_rate", "sp_mean_plies",
            "sp_mean_entropy", "tr_loss", "tr_value_mae", "t_selfplay", "t_learn"]
    rows = ["| " + " | ".join(k.replace("sp_", "").replace("tr_", "") for k in keys) + " |",
            "|" + "---|" * len(keys)]
    for r in iters[-25:]:
        rows.append("| " + " | ".join(
            (f"{r.get(k):.3f}" if isinstance(r.get(k), float) else str(r.get(k, "")))
            for k in keys) + " |")
    return "\n".join(rows)


# ----------------------------------------------------------------- traces
def list_checkpoints():
    out = []
    for d in RUNS.iterdir():
        if d.is_dir():
            out += [str(p.relative_to(ROOT)) for p in sorted(d.glob("*.pt"))]
    return out


def run_trace(ckpt: str, fen: str, steps: int, adaptive: bool):
    if not ckpt:
        return "Select a checkpoint.", "{}"
    cmd = [PY, "inspect_trace.py", "--checkpoint", ckpt, "--fen", fen or
           "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1",
           "--steps", str(steps), "--format", "both"]
    if adaptive:
        cmd.append("--adaptive")
    r = subprocess.run(cmd, cwd=ROOT, capture_output=True, text=True, timeout=300)
    return (r.stdout or "") + (("\n[stderr]\n" + r.stderr) if r.stderr else ""), ""


# ----------------------------------------------------------------- GPU smoke test
def gpu_smoke_test():
    """Tiny GPU sanity check. NOT used for training (5-min cap)."""
    import torch
    from tinychess.config import baseline_config
    from tinychess.model import TinyChess
    lines = [f"torch {torch.__version__}", f"cuda available: {torch.cuda.is_available()}"]
    if torch.cuda.is_available():
        dev = torch.device("cuda")
        lines.append(f"device: {torch.cuda.get_device_name(0)}")
        m = TinyChess(baseline_config()).to(dev)
        rep = m.param_report()
        lines.append(f"model params: {rep['total']:,}")
        sq = torch.randint(0, 13, (8, 64), device=dev)
        ex = torch.zeros(8, 64, 4, device=dev)
        gl = torch.zeros(8, 20, device=dev)
        cf = torch.zeros(8, 20, 7, dtype=torch.long, device=dev)
        cm = torch.ones(8, 20, dtype=torch.bool, device=dev)
        import time as _t
        t = _t.time()
        with torch.no_grad():
            out = m(sq, ex, gl, cf, cm, steps=8)
        lines.append(f"forward(8 steps, B=8): {(_t.time()-t)*1000:.1f} ms")
        lines.append(f"logits {tuple(out['logits'].shape)}  value {tuple(out['value'].shape)}")
        lines.append("GPU OK - but training stays on CPU by design.")
    else:
        lines.append("no CUDA visible")
    return "\n".join(lines)


ALLOWED_CMDS = {"train_selfplay.py", "inspect_trace.py", "tests/test_phase1.py",
                "scripts/bench_threads.py", "experiments/exp_ablation.py",
                "experiments/exp_learning.py", "experiments/exp_compute_scaling.py",
                "experiments/exp_faithfulness.py", "experiments/exp_precision.py",
                "experiments/analyse_grid.py", "experiments/exp_step_embedding.py",
                "experiments/exp_halting_routing.py",
                "experiments/exp_mate_capability.py",
                "experiments/exp_check_feature.py",
                "experiments/exp_mate_readout.py",
                "experiments/exp_check_intervention.py",
                "experiments/exp_endgame_conversion.py",
                "scripts/bench_parallel.py",
                "scripts/bench_learn.py",
                "experiments/exp_greedy_health.py"}


def run_job(name: str, script: str, argstr: str):
    """Launch a whitelisted research script as a detached background job."""
    script = script.strip()
    if script not in ALLOWED_CMDS:
        return f"script '{script}' not in the allowed list: {sorted(ALLOWED_CMDS)}"
    import shlex
    args = shlex.split(argstr or "")
    # The Space container is EPHEMERAL: a rebuild deletes runs/. Force periodic
    # pushes to the HF model repo so no experiment is ever lost.
    if script == "train_selfplay.py" and "--push-every" not in args:
        args += ["--push-every", "5"]
    cmd = [PY, script] + args
    return launch(name.strip() or "job", cmd)


def read_file_tail(relpath: str, n: int = 200):
    p = (ROOT / relpath.strip()).resolve()
    if not str(p).startswith(str(ROOT)) or not p.exists():
        return f"(no such file: {relpath})"
    if p.suffix in (".pt", ".pkl"):
        return f"{relpath}: binary, {p.stat().st_size/1e6:.2f} MB"
    lines = p.read_text(errors="replace").splitlines()
    return "\n".join(lines[-int(n):])


def list_files(pattern: str = "runs/**/*"):
    import glob
    out = []
    for f in sorted(glob.glob(str(ROOT / pattern), recursive=True))[:400]:
        p = Path(f)
        if p.is_file():
            out.append(f"{p.relative_to(ROOT)}  {p.stat().st_size/1024:.1f} KB")
    return "\n".join(out) or "(none)"


def env_audit():
    """Report the ACTUAL hardware of this Space (the research machine)."""
    import multiprocessing, platform, subprocess, torch
    L = []
    L.append(f"python {platform.python_version()}  torch {torch.__version__}")
    L.append(f"os.cpu_count(): {os.cpu_count()}")
    try:
        L.append(f"sched_getaffinity: {len(os.sched_getaffinity(0))}")
    except Exception:
        pass
    L.append(f"torch threads: {torch.get_num_threads()}")
    try:
        out = subprocess.run(["lscpu"], capture_output=True, text=True, timeout=20).stdout
        for k in ("Model name", "CPU(s):", "Thread(s) per core", "Core(s) per socket",
                  "Socket(s)", "L3 cache"):
            for line in out.splitlines():
                if line.strip().startswith(k):
                    L.append(line.strip()); break
    except Exception as e:
        L.append(f"lscpu failed: {e}")
    try:
        mt = open("/proc/meminfo").readline().split()
        L.append(f"MemTotal: {int(mt[1])/1024/1024:.1f} GB")
    except Exception:
        pass
    try:
        import shutil
        du = shutil.disk_usage("/data" if os.path.isdir("/data") else ".")
        L.append(f"disk free: {du.free/1e9:.1f} GB  (persistent /data: {os.path.isdir('/data')})")
    except Exception:
        pass
    L.append(f"cuda available: {torch.cuda.is_available()}")
    L.append(f"BLAS: {torch.__config__.show().split('BLAS_INFO=')[-1].split(',')[0] if 'BLAS_INFO' in torch.__config__.show() else '?'}")
    L.append(f"cpu capability: {torch.backends.cpu.get_cpu_capability()}")
    return "\n".join(L)


# ----------------------------------------------------------------- UI
with gr.Blocks(title="TinyChess Research Console") as demo:
    gr.Markdown("# TinyChess — ~100K-parameter self-play research console\n"
                "Training runs as detached CLI jobs; this UI only controls and inspects them.")

    with gr.Tab("Train"):
        with gr.Row():
            cfg_name = gr.Dropdown(["baseline", "memory", "refine", "halting",
                                    "router", "full", "mlp", "transformer"],
                                   value="baseline", label="config")
            job_name = gr.Textbox("sp_baseline", label="job / run name")
        with gr.Row():
            iters = gr.Number(200, label="iterations", precision=0)
            games = gr.Number(64, label="games / iter", precision=0)
            steps = gr.Number(4, label="recurrent steps", precision=0)
            threads = gr.Number(4, label="torch threads", precision=0)
        with gr.Row():
            updates = gr.Number(24, label="grad steps / iter", precision=0)
            batch = gr.Number(256, label="batch", precision=0)
            lr = gr.Number(3e-4, label="lr")
            seed = gr.Number(0, label="seed", precision=0)
        with gr.Row():
            b_start = gr.Button("Start", variant="primary")
            b_resume = gr.Button("Resume")
            b_stop = gr.Button("Stop", variant="stop")
        train_out = gr.Textbox(label="result", lines=4)

        def _start(cfg, name, it, g, st, th, up, ba, lr_, sd, resume):
            out = f"runs/{name}"
            cmd = [PY, "train_selfplay.py", "--config", cfg, "--out", out,
                   "--iters", str(int(it)), "--games", str(int(g)),
                   "--steps", str(int(st)), "--threads", str(int(th)),
                   "--updates", str(int(up)), "--batch", str(int(ba)),
                   "--lr", str(lr_), "--seed", str(int(sd))]
            if resume:
                cmd += ["--resume", f"{out}/latest.pt"]
            return launch(name, cmd)

        b_start.click(lambda *a: _start(*a, False),
                      [cfg_name, job_name, iters, games, steps, threads, updates,
                       batch, lr, seed], train_out)
        b_resume.click(lambda *a: _start(*a, True),
                       [cfg_name, job_name, iters, games, steps, threads, updates,
                        batch, lr, seed], train_out)
        b_stop.click(stop_job, job_name, train_out)

    with gr.Tab("Jobs"):
        jt = gr.Markdown()
        logname = gr.Textbox("sp_baseline", label="job")
        logbox = gr.Textbox(label="log tail", lines=25, max_lines=25)
        with gr.Row():
            gr.Button("Refresh jobs").click(job_table, None, jt)
            gr.Button("Refresh log").click(lambda n: tail_log(n, 80), logname, logbox)
        demo.load(job_table, None, jt)

    with gr.Tab("Metrics"):
        run_dd = gr.Dropdown(list_runs(), label="run")
        with gr.Row():
            gr.Button("Reload runs").click(lambda: gr.update(choices=list_runs()), None, run_dd)
            gr.Button("Plot").click(metrics_plot, run_dd, [gr.Plot(), gr.Markdown()])
        tbl = gr.Markdown()
        gr.Button("Table").click(summary_table, run_dd, tbl)

    with gr.Tab("Inspect trace"):
        ck = gr.Dropdown(list_checkpoints(), label="checkpoint")
        fen_in = gr.Textbox("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1",
                            label="FEN")
        with gr.Row():
            tsteps = gr.Number(8, label="steps", precision=0)
            tad = gr.Checkbox(False, label="adaptive halting")
        gr.Button("Reload checkpoints").click(
            lambda: gr.update(choices=list_checkpoints()), None, ck)
        trace_out = gr.Textbox(label="trace", lines=30)
        gr.Button("Run trace", variant="primary").click(
            run_trace, [ck, fen_in, tsteps, tad], [trace_out, gr.Textbox(visible=False)])

    with gr.Tab("Run job"):
        gr.Markdown("Launch a research script as a detached background job on "
                    "the Space CPUs. Jobs survive browser disconnects.")
        rj_name = gr.Textbox("job1", label="job name")
        rj_script = gr.Textbox("train_selfplay.py", label="script")
        rj_args = gr.Textbox("--config baseline --iters 10 --games 32 --threads 8 --out runs/test",
                             label="arguments")
        rj_out = gr.Textbox(label="result", lines=4)
        gr.Button("Launch", variant="primary").click(run_job,
                                                     [rj_name, rj_script, rj_args], rj_out)
        gr.Markdown("### Inspect files")
        f_pat = gr.Textbox("runs/**/*", label="glob")
        f_list = gr.Textbox(label="files", lines=12)
        gr.Button("List files").click(list_files, f_pat, f_list)
        f_path = gr.Textbox("runs/test/metrics.jsonl", label="file")
        f_n = gr.Number(100, label="tail lines", precision=0)
        f_body = gr.Textbox(label="contents", lines=20)
        gr.Button("Read file").click(read_file_tail, [f_path, f_n], f_body)

    with gr.Tab("Environment"):
        gr.Markdown("Actual hardware of this Space - the real research machine.")
        env_out = gr.Textbox(label="environment", lines=18)
        gr.Button("Audit environment", variant="primary").click(env_audit, None, env_out)

    with gr.Tab("GPU check"):
        gr.Markdown("ZeroGPU smoke test only. Training runs on CPU as detached "
                    "jobs because the GPU allocation is capped at ~5 minutes.")
        gpu_out = gr.Textbox(label="gpu", lines=10)
        gr.Button("Run GPU smoke test").click(gpu_smoke_test, None, gpu_out)

    with gr.Tab("Evaluate"):
        ck2 = gr.Dropdown(list_checkpoints(), label="checkpoint")
        init2 = gr.Textbox("", label="random-init checkpoint (optional)")
        egames = gr.Number(60, label="games per match", precision=0)
        ejob = gr.Textbox("eval1", label="job name")
        eout = gr.Textbox(label="result", lines=4)
        gr.Button("Reload").click(lambda: gr.update(choices=list_checkpoints()), None, ck2)
        gr.Button("Run evaluation", variant="primary").click(
            lambda c, i, g, n: launch(n, [PY, "experiments/exp_learning.py",
                                          "--checkpoint", c, "--games", str(int(g))]
                                      + (["--init", i] if i else [])),
            [ck2, init2, egames, ejob], eout)

if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0",
                        server_port=int(os.environ.get("PORT", 7860)),
                        ssr_mode=False, show_api=False)