"""
Live training dashboard: loss, validation error, learning rate, GPU, and what the model hears.

    python dashboard.py                  # opens http://localhost:8008

Leave it running in a second window next to train.py; the page updates itself every few seconds.
It reads runs/*/log.csv, status.json and samples.json (written by train.py) and samples the GPU
with nvidia-smi. Its Pause button asks train.py to save and exit (it leaves a file named "stop" in
the run's folder, which train.py checks every few seconds). Plain Python, nothing to install (CPU and RAM show too if psutil is installed).
To watch from your phone on the same Wi-Fi: python dashboard.py --host 0.0.0.0, then open
http://<this PC's IP>:8008.
"""
import argparse
import collections
import csv
import json
import math
import random
import os
import subprocess
import sys
import threading
import time
import webbrowser
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import parse_qs, unquote, urlparse

HERE = os.path.dirname(os.path.abspath(__file__))
PAGE = os.path.join(HERE, "dashboard.html")
FONT = os.path.normpath(os.path.join(HERE, "..", "app", "src", "main", "res", "font", "hafs.ttf"))
GPU_QUERY = "utilization.gpu,memory.used,memory.total,temperature.gpu,power.draw,name"

try:
    import psutil
except ImportError:
    psutil = None


class Monitor:
    """Samples the GPU (and CPU/RAM when psutil is available) every couple of seconds, keeping an hour."""

    def __init__(self, every=2.0, keep=1800):
        self.every = every
        self.samples = collections.deque(maxlen=keep)
        self.gpu_name = None
        self.nvidia_smi = True
        self.lock = threading.Lock()
        threading.Thread(target=self._loop, daemon=True).start()

    def _gpu(self):
        if not self.nvidia_smi:
            return None
        try:
            out = subprocess.run(["nvidia-smi", f"--query-gpu={GPU_QUERY}", "--format=csv,noheader,nounits"],
                                 capture_output=True, text=True, timeout=10,
                                 creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0))
        except FileNotFoundError:
            self.nvidia_smi = False  # no NVIDIA driver: stop asking
            return None
        except (OSError, subprocess.TimeoutExpired):
            return None
        line = out.stdout.strip().splitlines()[0] if out.returncode == 0 and out.stdout.strip() else ""
        parts = [p.strip() for p in line.split(",")]
        if len(parts) < 6:
            return None
        self.gpu_name = parts[5]

        def num(s):
            try:
                return float(s)
            except ValueError:
                return None  # "[N/A]" on some laptops

        return {"util": num(parts[0]), "mem": num(parts[1]), "mem_total": num(parts[2]),
                "temp": num(parts[3]), "power": num(parts[4])}

    def _loop(self):
        if psutil:
            psutil.cpu_percent(None)
        while True:
            started = time.time()
            sample = {"t": round(started, 2)}
            gpu = self._gpu()
            if gpu:
                sample.update(gpu)
            if psutil:
                sample["cpu"] = psutil.cpu_percent(None)
                sample["ram"] = round(psutil.virtual_memory().used / 2**30, 2)
            if len(sample) > 1:
                with self.lock:
                    self.samples.append(sample)
            time.sleep(max(0.2, self.every - (time.time() - started)))

    def since(self, t):
        with self.lock:
            rows = [s for s in self.samples if s["t"] > t]
        ram_total = round(psutil.virtual_memory().total / 2**30, 1) if psutil else None
        return {"name": self.gpu_name, "available": self.nvidia_smi, "ram_total": ram_total, "samples": rows}


def number(value):
    if value in ("", None):
        return None
    try:
        v = float(value)
    except ValueError:
        return None
    return v if math.isfinite(v) else None


def read_log(path):
    """log.csv as columns of numbers. Rows a resumed run redid (steps going backwards) are dropped."""
    rows = []
    with open(path, newline="", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        columns = reader.fieldnames or []
        for raw in reader:
            row = {k: number(raw.get(k)) for k in columns}
            if row.get("step") is None:
                continue  # a line cut short by a crash
            if rows and row["step"] < rows[-1]["step"]:
                while rows and rows[-1]["step"] >= row["step"]:
                    rows.pop()
            rows.append(row)
    return {k: [r[k] for r in rows] for k in columns}


def read_json(path):
    try:
        with open(path, encoding="utf-8") as f:
            return json.load(f)
    except (OSError, ValueError):
        return None


class Runs:
    """The runs folder, re-read only when a file changes."""

    FILES = ("log.csv", "status.json", "samples.json")

    def __init__(self, root):
        self.root = root
        self.cache = {}

    def names(self):
        if not os.path.isdir(self.root):
            return []
        return sorted(n for n in os.listdir(self.root)
                      if any(os.path.exists(os.path.join(self.root, n, f)) for f in self.FILES))

    def stop_file(self, name):
        return os.path.join(self.root, name, "stop")

    def request_pause(self, name, cancel=False):
        if name not in self.names():
            return False
        path = self.stop_file(name)
        if cancel:
            if os.path.exists(path):
                os.remove(path)
        else:
            with open(path, "w", encoding="utf-8") as f:
                f.write("pause requested from the dashboard\n")
        return True

    def _stamp(self, name):
        stamps = []
        for f in self.FILES:
            try:
                st = os.stat(os.path.join(self.root, name, f))
                stamps.append(f"{st.st_mtime_ns}-{st.st_size}")
            except OSError:
                stamps.append("-")
        return "_".join(stamps)

    def summary(self):
        out = []
        for name in self.names():
            status = read_json(os.path.join(self.root, name, "status.json")) or {}
            mtimes = [os.path.getmtime(os.path.join(self.root, name, f)) for f in self.FILES
                      if os.path.exists(os.path.join(self.root, name, f))]
            out.append({"name": name, "phase": status.get("phase"), "updated_at": max(mtimes, default=0),
                        "step": status.get("step"), "total_steps": status.get("total_steps"),
                        "version": self._stamp(name), "pause_requested": os.path.exists(self.stop_file(name))})
        return out

    def samples(self, name, seed, count=12):
        """
        `count` validation ayahs drawn with `seed` (Warsh first, then one reciter at a time, 3-20 s
        long) and what the model heard on them at every validation, from runs/<name>/heard/.
        Runs started before train.py saved every clip only have their fixed samples.json.
        """
        folder = os.path.join(self.root, name)
        heard_dir = os.path.join(folder, "heard")
        clips = read_json(os.path.join(heard_dir, "clips.json"))
        steps = sorted(f for f in os.listdir(heard_dir) if f.startswith("step_")) if clips else []
        if not steps:
            fixed = read_json(os.path.join(folder, "samples.json"))
            return dict(fixed or {"clips": [], "evals": []}, reshuffle=False)
        rnd = random.Random(seed)
        by_reciter = {}
        for i, c in enumerate(clips):
            if 3 <= c["duration"] <= 20:
                by_reciter.setdefault((c["riwaya"] != "warsh", c["reciter"]), []).append(i)
        queues = [rnd.sample(ids, len(ids)) for _, ids in sorted(by_reciter.items())]
        warsh = [q for (not_warsh, _), q in zip(sorted(by_reciter), queues) if not not_warsh]
        hafs = [q for (not_warsh, _), q in zip(sorted(by_reciter), queues) if not_warsh]
        rnd.shuffle(warsh)
        rnd.shuffle(hafs)
        picked = []
        while len(picked) < count and any(warsh + hafs):
            for q in warsh + hafs:
                if q and len(picked) < count:
                    picked.append(q.pop())
        evals = []
        for f in steps:
            key = (name, f)
            path = os.path.join(heard_dir, f)
            stamp = os.path.getmtime(path)
            if self.cache.get(key, (None,))[0] != stamp:
                data = read_json(path)
                if data is None:
                    continue
                self.cache[key] = (stamp, data)
            data = self.cache[key][1]
            evals.append({"step": data["step"], "hyps": [data["hyps"][i] if i < len(data["hyps"]) else "" for i in picked]})
        return {"clips": [clips[i] for i in picked], "evals": evals, "reshuffle": True}

    def run(self, name, known_version=None):
        if name not in self.names():
            return None
        version = self._stamp(name)
        if version == known_version:
            return {"version": version, "unchanged": True}
        cached = self.cache.get(name)
        if cached and cached["version"] == version:
            return cached
        folder = os.path.join(self.root, name)
        log_path = os.path.join(folder, "log.csv")
        try:
            log = read_log(log_path) if os.path.exists(log_path) else {}
        except OSError:
            log = {}
        result = {"version": version, "name": name, "log": log,
                  "status": read_json(os.path.join(folder, "status.json")),
                  "samples": read_json(os.path.join(folder, "samples.json"))}
        self.cache[name] = result
        return result


def make_handler(runs, monitor):
    class Handler(BaseHTTPRequestHandler):
        def log_message(self, *args):
            pass  # keep the console quiet

        def _send(self, body, content_type, status=200, cache=False):
            self.send_response(status)
            self.send_header("Content-Type", content_type)
            self.send_header("Content-Length", str(len(body)))
            self.send_header("Cache-Control", "max-age=86400" if cache else "no-store")
            self.end_headers()
            self.wfile.write(body)

        def _json(self, value, status=200):
            self._send(json.dumps(value, ensure_ascii=False, allow_nan=False).encode("utf-8"),
                       "application/json; charset=utf-8", status)

        def do_GET(self):
            url = urlparse(self.path)
            query = parse_qs(url.query)
            try:
                if url.path in ("/", "/index.html"):
                    with open(PAGE, "rb") as f:  # read each time: edits show on refresh
                        self._send(f.read(), "text/html; charset=utf-8")
                elif url.path == "/font/hafs.ttf" and os.path.exists(FONT):
                    with open(FONT, "rb") as f:
                        self._send(f.read(), "font/ttf", cache=True)
                elif url.path == "/api/runs":
                    self._json({"runs": runs.summary(), "now": time.time(), "root": os.path.abspath(runs.root)})
                elif url.path.startswith("/api/run/") and url.path.endswith("/samples"):
                    name = unquote(url.path[len("/api/run/"):-len("/samples")])
                    if name in runs.names():
                        self._json(runs.samples(name, int(query.get("seed", ["0"])[0])))
                    else:
                        self._json({"error": "no such run"}, 404)
                elif url.path.startswith("/api/run/"):
                    result = runs.run(unquote(url.path[len("/api/run/"):]), query.get("v", [None])[0])
                    self._json(result if result else {"error": "no such run"}, 200 if result else 404)
                elif url.path == "/api/system":
                    self._json(monitor.since(float(query.get("since", ["0"])[0])))
                else:
                    self._send(b"not found", "text/plain", 404)
            except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError):
                pass

        def do_POST(self):
            url = urlparse(self.path)
            parts = [unquote(p) for p in url.path.strip("/").split("/")]
            if len(parts) == 4 and parts[:2] == ["api", "run"] and parts[3] == "pause":
                cancel = parse_qs(url.query).get("cancel", ["0"])[0] == "1"
                ok = runs.request_pause(parts[2], cancel)
                self._json({"ok": ok}, 200 if ok else 404)
            else:
                self._send(b"not found", "text/plain", 404)

    return Handler


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--runs", default=os.path.join(HERE, "runs"), help="folder holding the runs")
    parser.add_argument("--port", type=int, default=8008)
    parser.add_argument("--host", default="127.0.0.1", help="0.0.0.0 to reach it from other devices")
    parser.add_argument("--no-browser", action="store_true", help="don't open the browser")
    args = parser.parse_args()

    server = ThreadingHTTPServer((args.host, args.port), make_handler(Runs(args.runs), Monitor()))
    server.daemon_threads = True
    url = f"http://localhost:{args.port}"
    print(f"dashboard: {url}  (runs from {os.path.abspath(args.runs)}; Ctrl+C to stop)")
    if not args.no_browser:
        threading.Timer(0.5, webbrowser.open, (url,)).start()
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    server.server_close()
    return 0


if __name__ == "__main__":
    sys.exit(main())