"""
Run registry and subprocess supervision for the UI.

Each run is a directory under output/<slug>/ containing everything needed to
interpret it later: the config that produced it, the solver log, the geometry
preview, the frames, and summary.json. That layout is the persistence
mechanism -- the registry file is a cache, and scan() can rebuild it from
directories alone. A UI that forgets your previous work because an index file
was lost is worse than no UI.

Progress is parsed from the solver's own stdout rather than tracked through a
side channel. The solver prints "running to t = T" once and "step N t=..."
periodically, which is enough for a percentage and an ETA, and it means the
UI cannot drift out of sync with what the solver is actually doing.
"""

import json
import os
import re
import signal
import subprocess
import sys
import threading
import time

OUTPUT_ROOT = "output"
REGISTRY = os.path.join(OUTPUT_ROOT, "_runs.json")

_TARGET_RE = re.compile(r"running to t = ([0-9.eE+-]+)")
_STEP_RE = re.compile(
    r"step\s+(\d+)\s+t=\s*([0-9.eE+-]+)\s+dt=([0-9.eE+-]+)\s+"
    r"div=([0-9.eE+-]+)\s+F=\(([^,]+),([^)]+)\)\s+([0-9.]+)s"
)
_ERROR_RE = re.compile(r"(Traceback|Error|ERROR|diverged|FloatingPointError)")


def slugify(name):
    base = os.path.splitext(os.path.basename(name))[0]
    s = re.sub(r"[^A-Za-z0-9._-]+", "-", base).strip("-").lower()
    return s or "run"


def unique_slug(name, root=OUTPUT_ROOT):
    slug = slugify(name)
    candidate = slug
    n = 2
    while os.path.exists(os.path.join(root, candidate)):
        candidate = f"{slug}-{n}"
        n += 1
    return candidate


class Run:
    def __init__(self, run_id, directory, label=""):
        self.id = run_id
        self.dir = directory
        self.label = label or run_id
        self.status = "queued"      # queued|preflight|running|done|failed|stopped
        self.created = time.time()
        self.finished = None
        self.proc = None
        self.log_path = os.path.join(directory, "run.log")
        self.error = None

    # -- persisted view ------------------------------------------------

    def to_dict(self):
        d = {
            "id": self.id,
            "label": self.label,
            "dir": self.dir.replace("\\", "/"),
            "status": self.status,
            "created": self.created,
            "finished": self.finished,
            "error": self.error,
        }
        d.update(self.progress())
        d["summary"] = self.summary()
        d["files"] = self.files()
        return d

    def summary(self):
        p = os.path.join(self.dir, "summary.json")
        if not os.path.exists(p):
            return None
        try:
            with open(p) as fh:
                return json.load(fh)
        except (json.JSONDecodeError, OSError):
            return None

    def files(self):
        """Everything produced, as browser-reachable paths."""
        out = []
        if not os.path.isdir(self.dir):
            return out
        for root, dirs, names in os.walk(self.dir):
            # Frame directories hold hundreds of files; summarise instead of
            # listing them all, or the UI payload balloons past a megabyte.
            rel_root = os.path.relpath(root, self.dir).replace("\\", "/")
            if rel_root.endswith("frames") or "/png" in rel_root:
                out.append({
                    "name": rel_root + "/",
                    "kind": "folder",
                    "count": len(names),
                    "path": None,
                })
                dirs[:] = [d for d in dirs if d != "png"]
                continue
            for n in sorted(names):
                full = os.path.join(root, n)
                rel = os.path.relpath(full, OUTPUT_ROOT).replace("\\", "/")
                try:
                    size = os.path.getsize(full)
                except OSError:
                    size = 0
                out.append({
                    "name": (rel_root + "/" + n) if rel_root != "." else n,
                    "kind": os.path.splitext(n)[1].lstrip(".").lower(),
                    "size": size,
                    "path": "/files/" + rel,
                })
        return out

    # -- live progress -------------------------------------------------

    def progress(self):
        info = {
            "percent": 0.0, "sim_time": 0.0, "target_time": None,
            "step": 0, "dt": None, "divergence": None,
            "fx": None, "fy": None, "elapsed": 0.0, "eta": None,
            "phase": self.status,
        }
        # A finished run reads 100% even with no log to parse -- runs started
        # from the terminal write their log elsewhere, and a completed project
        # showing an empty progress bar looks like a failure.
        if self.status == "done":
            info["percent"] = 100.0
        if not os.path.exists(self.log_path):
            return info
        try:
            with open(self.log_path, "r", errors="replace") as fh:
                text = fh.read()
        except OSError:
            return info

        m = _TARGET_RE.search(text)
        if m:
            info["target_time"] = float(m.group(1))

        last = None
        for last in _STEP_RE.finditer(text):
            pass
        if last:
            info["step"] = int(last.group(1))
            info["sim_time"] = float(last.group(2))
            info["dt"] = float(last.group(3))
            info["divergence"] = float(last.group(4))
            info["fx"] = float(last.group(5))
            info["fy"] = float(last.group(6))
            info["elapsed"] = float(last.group(7))
            if info["target_time"]:
                frac = min(info["sim_time"] / info["target_time"], 1.0)
                info["percent"] = 100.0 * frac
                if frac > 0.01:
                    info["eta"] = info["elapsed"] * (1.0 / frac - 1.0)
        elif info["target_time"] is None and self.status == "running":
            info["phase"] = "rasterising geometry"

        if self.status == "done":
            info["percent"] = 100.0
        return info

    def log_tail(self, lines=40):
        if not os.path.exists(self.log_path):
            return ""
        try:
            with open(self.log_path, "r", errors="replace") as fh:
                return "".join(fh.readlines()[-lines:])
        except OSError:
            return ""


class RunManager:
    """Owns the run registry and the subprocesses."""

    def __init__(self, root=OUTPUT_ROOT):
        self.root = root
        self.runs = {}
        self._lock = threading.Lock()
        os.makedirs(root, exist_ok=True)
        self.scan()

    # -- discovery -----------------------------------------------------

    def scan(self):
        """Rebuild the registry from what is actually on disk.

        Directory contents are the source of truth, not the registry file, so
        a corrupted or deleted index never loses a completed project.
        """
        known = {}
        if os.path.exists(REGISTRY):
            try:
                with open(REGISTRY) as fh:
                    known = {r["id"]: r for r in json.load(fh)}
            except (json.JSONDecodeError, OSError, KeyError, TypeError):
                known = {}

        with self._lock:
            for name in sorted(os.listdir(self.root)):
                d = os.path.join(self.root, name)
                if not os.path.isdir(d) or name.startswith("_"):
                    continue
                if name in self.runs:
                    continue
                run = Run(name, d, label=known.get(name, {}).get("label", name))
                run.created = known.get(name, {}).get(
                    "created", os.path.getmtime(d)
                )
                if os.path.exists(os.path.join(d, "summary.json")):
                    run.status = "done"
                elif os.path.exists(os.path.join(d, "run.log")):
                    # A log with no summary and no live process means the run
                    # died -- most likely the app was closed mid-run.
                    run.status = known.get(name, {}).get("status", "failed")
                    if run.status in ("running", "queued", "preflight"):
                        run.status = "failed"
                        run.error = "interrupted (process no longer running)"
                else:
                    # Pre-flight-only runs leave a config and a preview but no
                    # log. "empty" reads like a fault; they are a deliberate
                    # and useful artefact.
                    run.status = "checked"
                self.runs[name] = run
        return self.list()

    def save(self):
        try:
            with open(REGISTRY, "w") as fh:
                json.dump(
                    [{"id": r.id, "label": r.label, "created": r.created,
                      "status": r.status} for r in self.runs.values()],
                    fh, indent=2,
                )
        except OSError:
            pass

    # -- lifecycle -----------------------------------------------------

    def create(self, label, config_text):
        run_id = unique_slug(label)
        d = os.path.join(self.root, run_id)
        os.makedirs(d, exist_ok=True)
        cfg_path = os.path.join(d, "config.yaml")
        with open(cfg_path, "w") as fh:
            fh.write(config_text)
        run = Run(run_id, d, label=label)
        with self._lock:
            self.runs[run_id] = run
        self.save()
        return run

    def start(self, run_id, extra_args=()):
        run = self.runs[run_id]
        cfg = os.path.join(run.dir, "config.yaml")
        cmd = [sys.executable, "-u", "waterfall.py", cfg, *extra_args]
        logf = open(run.log_path, "w", buffering=1, errors="replace")
        # CREATE_NEW_PROCESS_GROUP lets stop() signal the child on Windows
        # without also killing the server process hosting the UI.
        flags = subprocess.CREATE_NEW_PROCESS_GROUP if os.name == "nt" else 0
        run.proc = subprocess.Popen(
            cmd, stdout=logf, stderr=subprocess.STDOUT,
            cwd=os.getcwd(), creationflags=flags,
        )
        run.status = "running"
        self.save()

        def watch():
            code = run.proc.wait()
            logf.close()
            run.finished = time.time()
            if run.status == "stopped":
                pass
            elif code == 0 and os.path.exists(
                os.path.join(run.dir, "summary.json")
            ):
                run.status = "done"
            else:
                run.status = "failed"
                tail = run.log_tail(30)
                m = _ERROR_RE.search(tail)
                run.error = (
                    tail.strip().splitlines()[-1] if m and tail.strip()
                    else f"exited with code {code}"
                )
            self.save()

        threading.Thread(target=watch, daemon=True).start()
        return run

    def stop(self, run_id):
        run = self.runs.get(run_id)
        if not run or not run.proc or run.proc.poll() is not None:
            return False
        run.status = "stopped"
        try:
            if os.name == "nt":
                run.proc.send_signal(signal.CTRL_BREAK_EVENT)
                time.sleep(1.0)
            run.proc.terminate()
        except (OSError, ValueError):
            pass
        self.save()
        return True

    def delete(self, run_id):
        import shutil
        run = self.runs.get(run_id)
        if not run:
            return False
        self.stop(run_id)
        shutil.rmtree(run.dir, ignore_errors=True)
        with self._lock:
            self.runs.pop(run_id, None)
        self.save()
        return True

    # -- views ---------------------------------------------------------

    def list(self):
        return sorted(
            (r.to_dict() for r in self.runs.values()),
            key=lambda r: r["created"], reverse=True,
        )

    def get(self, run_id):
        run = self.runs.get(run_id)
        if not run:
            return None
        d = run.to_dict()
        d["log"] = run.log_tail(200)
        return d

    def active(self):
        return [r for r in self.runs.values()
                if r.status in ("running", "queued", "preflight")]