"""Streaming EF5 runner for the live dashboard.

Runs the fork's ef5 non-blocking and streams outputs as they are written:
  - ts.<gauge>.<model>.csv  -> hydrograph rows (fork flushes every timestep)
  - q.<datetime>.<model>.tif -> 2-D streamflow frames (OUTPUT_GRIDS=STREAMFLOW)

`stream_run()` yields event dicts the Gradio app consumes. `MockEF5` writes the
same two output streams so the UI + reader can be tested without the binary.
"""
from __future__ import annotations

import os
import re
import glob
import math
import time
import subprocess
import threading
from dataclasses import dataclass, field
from datetime import datetime, timedelta

TS_TIME_FMT = "%Y-%m-%d %H:%M"

# Watchdog policy: a run that KEEPS PRODUCING OUTPUT is never killed for
# running long. Killed only when it goes silent (no new ts rows / q grids)
# for STALL_TIMEOUT_S, or exceeds the RUN_TIMEOUT_S hard cap.
RUN_TIMEOUT_S = float(os.environ.get("CREST_RUN_TIMEOUT_S", str(10 * 3600)))    # 10 h
STALL_TIMEOUT_S = float(os.environ.get("CREST_STALL_TIMEOUT_S", str(45 * 60)))  # 45 min


# --------------------------------------------------------------------------- #
# output readers
# --------------------------------------------------------------------------- #
# extra ts.csv columns streamed for the click-a-timestep readout (header token
# -> row key); real EF5 appends Temp/SWE when SNOW17 is on
_EXTRA_COLS = (("pet", "pet"), ("sm", "sm"), ("groundwater", "gw"),
               ("fast flow", "fast"), ("slow flow", "slow"), ("base flow", "base"),
               ("temp", "temp"), ("swe", "swe"))


def _col_map(header: list[str]) -> dict[str, int]:
    m: dict[str, int] = {}
    for i, h in enumerate(header):
        hl = h.strip().lower()
        for tok, key in _EXTRA_COLS:
            if hl.startswith(tok) and key not in m:
                m[key] = i
    return m


def _parse_ts_line(header: list[str], line: str, colmap: dict | None = None) -> dict | None:
    parts = line.rstrip("\n").split(",")
    if len(parts) < 4 or parts[0] == "Time":
        return None
    def f(x):
        try:
            v = float(x)
        except ValueError:
            return None
        return None if (math.isnan(v) or math.isinf(v)) else v
    row = {"time": parts[0], "sim_q": f(parts[1]), "obs_q": f(parts[2]), "precip": f(parts[3])}
    for key, i in (colmap or {}).items():
        if i < len(parts):
            v = f(parts[i])
            if v is not None:
                row[key] = v
    return row


class _TsTail:
    """Incrementally read new rows from a ts.csv as it grows."""
    def __init__(self, path: str):
        self.path = path
        self.header: list[str] | None = None
        self.colmap: dict[str, int] = {}
        self.pos = 0

    def read_new(self) -> list[dict]:
        if not os.path.exists(self.path):
            return []
        with open(self.path, "rb") as fh:          # binary: exact byte position
            fh.seek(self.pos)
            data = fh.read()
        nl = data.rfind(b"\n")                       # consume complete lines only
        if nl == -1:
            return []
        self.pos += nl + 1
        rows = []
        for raw in data[:nl + 1].decode("utf-8", "replace").splitlines():
            if self.header is None and raw.startswith("Time"):
                self.header = raw.split(",")
                self.colmap = _col_map(self.header)
                continue
            r = _parse_ts_line(self.header or [], raw, self.colmap)
            if r:
                rows.append(r)
        return rows


def _new_q_grids(output_dir: str, model: str, seen: set) -> list[str]:
    found = sorted(glob.glob(os.path.join(output_dir, f"q.*.{model}.tif")),
                   key=lambda p: os.path.getmtime(p))
    fresh = [p for p in found if p not in seen]
    seen.update(fresh)
    return fresh


# --------------------------------------------------------------------------- #
# streaming driver
# --------------------------------------------------------------------------- #
@dataclass
class RunHandle:
    output_dir: str
    ts_path: str
    model: str = "crestphys"
    proc: subprocess.Popen | None = None
    stop_flag: threading.Event = field(default_factory=threading.Event)

    def alive(self) -> bool:
        if self.proc is not None:
            return self.proc.poll() is None
        return not self.stop_flag.is_set()

    def kill(self):
        """Hard-stop a stuck run (watchdog)."""
        if self.proc is not None and self.proc.poll() is None:
            try:
                self.proc.kill()
                self.proc.wait(timeout=10)
            except Exception:
                pass
        self.stop_flag.set()


def run_ef5(control_path: str, output_dir: str, gauge_id: str, model: str = "crestphys",
            ef5_bin: str = "./EF5/bin/ef5") -> RunHandle:
    os.makedirs(output_dir, exist_ok=True)
    log = open(os.path.join(output_dir, "ef5_run.log"), "w")
    cmd = [ef5_bin, control_path]
    if os.name != "nt":            # line-buffer stdout so crashes keep the trail
        cmd = ["stdbuf", "-oL", "-eL"] + cmd
    proc = subprocess.Popen(cmd, stdout=log, stderr=subprocess.STDOUT, text=True)
    # EF5 lowercases the gauge id in the timeseries filename. All-digit USGS ids
    # are unaffected, but a virtual point's "V…" id is written as "v…", so the
    # expected path must lowercase too or the rows are never read (0-row runs).
    ts = os.path.join(output_dir, f"ts.{gauge_id.lower()}.{model}.csv")
    return RunHandle(output_dir=output_dir, ts_path=ts, model=model, proc=proc)


def stream_run(handle: RunHandle, poll: float = 0.4, timeout: float = RUN_TIMEOUT_S,
               stall_timeout: float = STALL_TIMEOUT_S, cancel=None):
    """Yield {'kind': 'hydro'|'q2d'|'done', ...} as EF5 writes output.

    A run that keeps producing output is left alone no matter how long it
    takes. The process is KILLED only when it goes silent for `stall_timeout`
    (stuck), passes the `timeout` hard cap, or `cancel` (a threading.Event)
    is set — the user stopped or superseded the job."""
    tail = _TsTail(handle.ts_path)
    seen: set = set()
    t0 = time.time()
    last_progress = t0                       # last time ANY output appeared
    while True:
        if cancel is not None and cancel.is_set():
            handle.kill()
            yield {"kind": "done", "returncode": -9, "cancelled": True}
            return
        rows = tail.read_new()
        if rows:
            yield {"kind": "hydro", "rows": rows}
        fresh = _new_q_grids(handle.output_dir, handle.model, seen)
        for p in fresh:
            yield {"kind": "q2d", "path": p}
        if rows or fresh:
            last_progress = time.time()
        done = not handle.alive()
        if done:
            # final drain
            rows = tail.read_new()
            if rows:
                yield {"kind": "hydro", "rows": rows}
            for p in _new_q_grids(handle.output_dir, handle.model, seen):
                yield {"kind": "q2d", "path": p}
            rc = handle.proc.returncode if handle.proc else 0
            yield {"kind": "done", "returncode": rc}
            return
        now = time.time()
        if now - last_progress > stall_timeout:
            handle.kill()                    # silent too long -> stuck
            yield {"kind": "done", "returncode": -1,
                   "error": f"no model output for {stall_timeout / 60:.0f} min — "
                            "killed (stall watchdog)"}
            return
        if now - t0 > timeout:
            handle.kill()                    # absolute cap
            yield {"kind": "done", "returncode": -1,
                   "error": f"killed after {timeout / 3600:.1f} h (run-time cap)"}
            return
        time.sleep(poll)


# --------------------------------------------------------------------------- #
# mock EF5 (for local testing without the binary)
# --------------------------------------------------------------------------- #
class MockEF5:
    """Writes ts.csv rows + q.*.tif grids incrementally, like a live EF5 run."""
    HEADER = ("Time,Discharge(m^3 s^-1),Observed(m^3 s^-1),Precip(mm h^-1),PET(mm h^-1),"
              "SM(%),Groundwater (mm),Fast Flow(mm*1000),Slow Flow(mm*1000),Base Flow(mm*1000)")

    def __init__(self, output_dir, gauge_id="01011000", model="crestphys",
                 bounds=(-69.83, 46.32, -68.33, 47.82), n_steps=48, t0=None,
                 dt_hours=1, delay=0.15, write_grids=True, facc_path=None,
                 has_obs=True):
        self.output_dir = output_dir
        self.gauge_id, self.model = gauge_id, model
        self.bounds, self.n_steps = bounds, n_steps
        self.t0 = t0 or datetime(2025, 7, 3, 0, 0)
        self.dt = timedelta(hours=dt_hours)
        self.delay, self.write_grids = delay, write_grids
        self.facc_path = facc_path            # clipped flow-acc -> realistic network
        self.has_obs = has_obs                # ungauged points: EF5 writes NaN obs
        self.ts_path = os.path.join(output_dir, f"ts.{gauge_id}.{model}.csv")

    def handle(self) -> RunHandle:
        return RunHandle(output_dir=self.output_dir, ts_path=self.ts_path, model=self.model)

    def start(self) -> RunHandle:
        os.makedirs(self.output_dir, exist_ok=True)
        h = self.handle()
        threading.Thread(target=self._run, args=(h,), daemon=True).start()
        return h

    def _run(self, h: RunHandle):
        import math
        import numpy as np
        try:
            with open(self.ts_path, "w") as ts:
                ts.write(self.HEADER + "\n"); ts.flush()
                for k in range(self.n_steps):
                    t = self.t0 + k * self.dt
                    # a plausible flood pulse
                    q = 5 + 60 * math.exp(-((k - self.n_steps * 0.4) ** 2) / (2 * (self.n_steps * 0.12) ** 2))
                    obs = q * (0.9 + 0.05 * math.sin(k / 3.0))
                    obs_s = f"{obs:.2f}" if self.has_obs else "nan"   # ungauged: no obs
                    p = max(0.0, 8 * math.exp(-((k - self.n_steps * 0.3) ** 2) / (2 * 3.0 ** 2)))
                    sm = 55 + 30 * (q - 5) / 60.0            # wets up with the pulse
                    gw = 7.0 + 0.02 * k
                    ts.write(f"{t:%Y-%m-%d %H:%M},{q:.2f},{obs_s},{p:.2f},0.10,"
                             f"{sm:.2f},{gw:.2f},0.0000,0.0000,0.6960\n")
                    ts.flush()
                    if self.write_grids:
                        self._grid(t, q, np)
                    time.sleep(self.delay)
        finally:
            h.stop_flag.set()          # signal completion to stream_run()

    def _grid(self, t: datetime, qpeak: float, np):
        """Write a per-timestep discharge grid. If a clipped flow-accumulation
        raster is available, put discharge on the actual stream network
        (discharge ~ upstream area), like a real EF5 q grid; else a fallback."""
        import rasterio
        from rasterio.transform import from_bounds as tr_from_bounds
        name = os.path.join(self.output_dir, f"q.{t:%Y%m%d%H%M}.{self.model}.tif")
        if self.facc_path and os.path.exists(self.facc_path):
            with rasterio.open(self.facc_path) as ds:
                scale = max(1, max(ds.width, ds.height) // 500)     # decimate
                oh, ow = max(1, ds.height // scale), max(1, ds.width // scale)
                facc = ds.read(1, out_shape=(oh, ow)).astype("float32")
                b, nod = ds.bounds, ds.nodata
            facc = np.where((facc == nod) | (facc <= 0), np.nan, facc)
            fmax = np.nanmax(facc)
            fmax = fmax if np.isfinite(fmax) and fmax > 0 else 1.0
            thr = np.nanpercentile(facc, 96)                        # top ~4% = channels
            onnet = np.isfinite(facc) & (facc >= thr)
            disch = qpeak * (facc / fmax)                          # ~ drainage area
            grid = np.where(onnet, disch, -9999.0).astype("float32")
            transform = tr_from_bounds(b.left, b.bottom, b.right, b.top, grid.shape[1], grid.shape[0])
        else:
            W, S, E, N = self.bounds
            nx = ny = 60
            yy, xx = np.mgrid[0:ny, 0:nx]
            chan = np.exp(-((xx - nx / 2) ** 2) / (2 * 3.0 ** 2))
            grid = np.where(chan * (yy / ny) > 0.05, qpeak * chan * (yy / ny), -9999.0).astype("float32")
            transform = tr_from_bounds(W, S, E, N, nx, ny)
        with rasterio.open(name, "w", driver="GTiff", height=grid.shape[0], width=grid.shape[1],
                           count=1, dtype="float32", crs="EPSG:4326", transform=transform,
                           nodata=-9999.0) as dst:
            dst.write(grid, 1)


if __name__ == "__main__":
    import sys, tempfile
    out = sys.argv[1] if len(sys.argv) > 1 else tempfile.mkdtemp()
    print("mock run ->", out)
    mock = MockEF5(out, n_steps=12, delay=0.1)
    h = mock.start()
    n_hydro = n_q = 0
    for ev in stream_run(h, poll=0.1):
        if ev["kind"] == "hydro":
            n_hydro += len(ev["rows"])
            last = ev["rows"][-1]
            print(f"  hydro row {last['time']} simQ={last['sim_q']:.1f} obsQ={last['obs_q']:.1f} P={last['precip']:.1f}")
        elif ev["kind"] == "q2d":
            n_q += 1
            print(f"  q2d frame -> {os.path.basename(ev['path'])}")
        elif ev["kind"] == "done":
            print(f"  DONE rc={ev['returncode']} | {n_hydro} hydro rows, {n_q} q-grids")