backend/flowtwin/runtime/session.py

"""Simulation sessions: the live runtime behind the dashboard.

A session owns one simulator, advances it on a wall-clock timer at the
requested speed multiplier, and publishes state frames to any connected
dashboards. Everything expensive (a step, a counterfactual sweep) runs off the
event loop so the WebSocket never stalls.

`ReplaySession` implements the same interface from a precomputed recording. It
exists so that a demo can continue if a live run cannot be created — see
`docs/DEMO.md`. It is never used unless the live path fails or is explicitly
requested.
"""

from __future__ import annotations

import asyncio
import json
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import numpy as np

from ..config import FALLBACK_DIR, Settings
from ..crowd.density import classify, level_name
from ..crowd.flow import build_alerts, detect_bottlenecks, primary_bottleneck
from ..prediction.inference import DensityPredictor
from ..simulation.agents import POLICY_ADAPTIVE, POLICY_BY_NAME, POLICY_SHORTEST
from ..simulation.engine import RunOverrides, Simulator
from ..strategy.engine import StrategyEngine
from ..venue import Scenario, Venue, compile_venue, load_scenario, load_venue

SPEED_CHOICES = (1, 2, 5, 10, 20, 40)


@dataclass
class SessionConfig:
    venue_id: str
    scenario_id: str
    seed: int
    crowd_size: int | None = None
    release_ramp_s: float | None = None
    compliance_scale: float = 1.0
    routing_policy: str = "shortest_path"
    capacity_overrides: dict[str, float] = field(default_factory=dict)
    event_factor_overrides: dict[str, float] = field(default_factory=dict)
    speed: int = 10
    autoplay: bool = False

    def as_dict(self) -> dict[str, Any]:
        return {
            "venue_id": self.venue_id,
            "scenario_id": self.scenario_id,
            "seed": self.seed,
            "crowd_size": self.crowd_size,
            "release_ramp_s": self.release_ramp_s,
            "compliance_scale": self.compliance_scale,
            "routing_policy": self.routing_policy,
            "capacity_overrides": dict(self.capacity_overrides),
            "event_factor_overrides": dict(self.event_factor_overrides),
            "speed": self.speed,
        }


class Broadcaster:
    """Fan-out of state frames to connected WebSocket clients."""

    def __init__(self) -> None:
        self._subscribers: set[asyncio.Queue] = set()

    def subscribe(self) -> asyncio.Queue:
        q: asyncio.Queue = asyncio.Queue(maxsize=4)
        self._subscribers.add(q)
        return q

    def unsubscribe(self, q: asyncio.Queue) -> None:
        self._subscribers.discard(q)

    @property
    def count(self) -> int:
        return len(self._subscribers)

    def publish(self, message: dict[str, Any]) -> None:
        for q in list(self._subscribers):
            if q.full():
                # Drop the oldest frame rather than block the simulation: a
                # slow client must not slow the venue down.
                try:
                    q.get_nowait()
                except asyncio.QueueEmpty:
                    pass
            try:
                q.put_nowait(message)
            except asyncio.QueueFull:
                pass


class SimulationSession:
    """A live, running simulation with its intelligence stack attached."""

    kind = "live"

    def __init__(self, config: SessionConfig, settings: Settings) -> None:
        self.id = uuid.uuid4().hex[:12]
        self.config = config
        self.settings = settings
        self.created_at = time.time()

        self.venue_model: Venue = load_venue(config.venue_id)
        self.compiled = compile_venue(config.venue_id)
        self.scenario: Scenario = load_scenario(config.scenario_id)
        if self.scenario.venue_id != config.venue_id:
            raise ValueError(
                f"scenario {config.scenario_id!r} belongs to venue "
                f"{self.scenario.venue_id!r}, not {config.venue_id!r}"
            )

        overrides = RunOverrides(
            crowd_size=config.crowd_size,
            release_ramp_s=config.release_ramp_s,
            compliance_scale=config.compliance_scale,
            routing_policy=POLICY_BY_NAME.get(config.routing_policy, POLICY_SHORTEST),
            capacity_overrides=dict(config.capacity_overrides),
            event_factor_overrides=dict(config.event_factor_overrides),
        )
        self.sim = Simulator(self.compiled, self.scenario, settings,
                             seed=config.seed, overrides=overrides)

        self.predictor = DensityPredictor(settings)
        self.strategy = StrategyEngine(settings, self.predictor)
        self.broadcaster = Broadcaster()

        self.speed = int(config.speed)
        self.playing = bool(config.autoplay)
        self.finished = False
        self.frame_index = 0
        self.last_error: str | None = None
        self.last_strategy_run: dict[str, Any] | None = None
        self._task: asyncio.Task | None = None
        self._lock = asyncio.Lock()
        self._busy = False
        self.last_seen = time.time()

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

    def start_loop(self) -> None:
        if self._task is None or self._task.done():
            self._task = asyncio.create_task(self._run_loop())

    async def close(self) -> None:
        self.playing = False
        if self._task is not None:
            self._task.cancel()
            try:
                await self._task
            except (asyncio.CancelledError, Exception):
                pass
            self._task = None

    async def _run_loop(self) -> None:
        interval = self.settings.server.frame_interval_s
        while True:
            started = time.perf_counter()
            # A session with nobody watching does no work. Without this, a
            # reloaded browser tab leaves an orphaned simulation stepping
            # forever and building frames no one reads, which starves the
            # event loop and makes new runs appear to hang.
            if self.broadcaster.count == 0:
                await asyncio.sleep(0.4)
                continue
            self.last_seen = time.time()
            if self.playing and not self.finished and not self._busy:
                sim_seconds = self.speed * interval
                steps = max(1, int(round(sim_seconds / self.sim.dt)))
                try:
                    await asyncio.to_thread(self._advance, steps)
                except Exception as exc:                       # pragma: no cover
                    self.last_error = f"{type(exc).__name__}: {exc}"
                    self.playing = False
                self.broadcaster.publish(self.frame())
            elapsed = time.perf_counter() - started
            await asyncio.sleep(max(0.01, interval - elapsed))

    def _advance(self, steps: int) -> None:
        for _ in range(steps):
            if self.sim.is_complete or self.sim.time >= self.scenario.duration_s:
                self.finished = True
                self.playing = False
                return
            self.sim.step()

    # -- controls ----------------------------------------------------------

    def play(self) -> None:
        if not self.finished:
            self.playing = True

    def pause(self) -> None:
        self.playing = False

    def set_speed(self, speed: int) -> None:
        self.speed = int(min(max(speed, 1), max(SPEED_CHOICES)))

    async def step_once(self, seconds: float = 10.0) -> None:
        steps = max(1, int(round(seconds / self.sim.dt)))
        await asyncio.to_thread(self._advance, steps)
        self.broadcaster.publish(self.frame())

    async def run_to(self, target_time_s: float) -> None:
        """Advance to a specific simulated time (used by the guided demo)."""
        steps = max(0, int(round((target_time_s - self.sim.time) / self.sim.dt)))
        if steps:
            await asyncio.to_thread(self._advance, steps)
        self.broadcaster.publish(self.frame())

    def trigger_event(self, index: int) -> dict[str, Any]:
        result = self.sim.trigger_event(index)
        self.broadcaster.publish(self.frame())
        return result

    # -- intelligence ------------------------------------------------------

    async def evaluate_strategies(self, horizon_s: float | None = None,
                                  strategy_ids: list[str] | None = None
                                  ) -> dict[str, Any]:
        async with self._lock:
            self._busy = True
            try:
                result = await asyncio.to_thread(
                    self.strategy.evaluate, self.sim, horizon_s, strategy_ids)
            finally:
                self._busy = False
        self.last_strategy_run = result
        self.broadcaster.publish({"type": "strategy", "session_id": self.id,
                                  "payload": result})
        return result

    async def apply_strategy(self, strategy_id: str) -> dict[str, Any]:
        async with self._lock:
            self._busy = True
            try:
                result = await asyncio.to_thread(self.strategy.apply, self.sim, strategy_id)
            finally:
                self._busy = False
        self.broadcaster.publish(self.frame())
        return result

    # -- serialisation -----------------------------------------------------

    def _edge_payload(self) -> list[dict[str, Any]]:
        """One entry per *physical* corridor, using the loaded direction."""
        v = self.compiled
        st = self.sim.state
        warning = self.venue_model.warning_density
        critical = self.venue_model.critical_density

        pair = v.pair_of
        has_pair = pair >= 0
        rev_in = np.zeros(v.n_edges)
        rev_in[has_pair] = st.edge_inflow_ppm[pair[has_pair]]
        dominant = st.edge_inflow_ppm >= rev_in

        levels = classify(st.edge_density, warning, critical)
        out: list[dict[str, Any]] = []
        seen: set[str] = set()
        for i in range(v.n_edges):
            base = v.edge_base_id[i]
            if base in seen or not dominant[i]:
                continue
            seen.add(base)
            out.append({
                "id": base,
                "dir": v.edge_ids[i],
                "reversed": bool(v.edge_reversed[i]),
                "d": round(float(st.edge_density[i]), 3),
                "dl": round(float(st.edge_peak_local_density[i]), 2),
                "v": round(float(st.edge_velocity[i]), 2),
                "in": round(float(st.edge_inflow_ppm[i])),
                "out": round(float(st.edge_outflow_ppm[i])),
                "q": int(st.edge_queue[i]),
                "occ": int(st.phys_occupancy[i]),
                "u": round(float(st.edge_inflow_ppm[i] / max(v.edge_capacity_ppm[i], 1)), 2),
                "g": round(float(st.edge_density_growth[i]), 3),
                "r": round(float(st.edge_risk[i]), 3),
                "lvl": level_name(int(levels[i])),
            })
        # Any corridor whose two directions are both idle still needs an entry.
        for i in range(v.n_edges):
            base = v.edge_base_id[i]
            if base in seen:
                continue
            seen.add(base)
            out.append({"id": base, "dir": v.edge_ids[i],
                        "reversed": bool(v.edge_reversed[i]),
                        "d": 0.0, "dl": 0.0, "v": round(self.settings.movement.free_speed_mps, 2),
                        "in": 0, "out": 0, "q": 0, "occ": 0, "u": 0.0, "g": 0.0,
                        "r": 0.0, "lvl": "clear"})
        return out

    def _node_payload(self) -> list[dict[str, Any]]:
        v = self.compiled
        st = self.sim.state
        levels = classify(st.node_density, self.venue_model.warning_density,
                          self.venue_model.critical_density)
        out = []
        for i, node in enumerate(self.venue_model.nodes):
            rate = float(v.node_service_ppm[i])
            mult = float(self.sim.node_budget.multiplier[i])
            out.append({
                "id": node.id,
                "occ": int(st.node_occupancy[i]),
                "d": round(float(st.node_density[i]), 3),
                "q": int(st.node_queue[i]),
                "thr": round(float(st.node_throughput_ppm[i])),
                "cap": None if not np.isfinite(rate) else round(rate * mult),
                "cap_base": None if not np.isfinite(rate) else round(rate),
                "cap_pct": round(100 * mult),
                "r": round(float(st.node_risk[i]), 3),
                "lvl": level_name(int(levels[i])),
            })
        return out

    def frame(self, include_agents: bool = True) -> dict[str, Any]:
        """One state frame for the dashboard."""
        sim = self.sim
        m = sim.metrics()
        preds = self.predictor.predict(sim)
        bottlenecks = detect_bottlenecks(sim, limit=6)
        alerts = build_alerts(sim, bottlenecks, preds)
        primary = primary_bottleneck(sim, preds)

        agents = (sim.agent_sample(self.settings.simulation.render_agent_budget)
                  if include_agents else {"x": [], "y": [], "v": [],
                                          "sampled": 0, "total": 0, "ratio": 1.0})

        self.frame_index += 1
        return {
            "type": "frame",
            "session_id": self.id,
            "kind": self.kind,
            "frame": self.frame_index,
            "t_s": round(sim.time, 1),
            "duration_s": self.scenario.duration_s,
            "playing": self.playing,
            "finished": self.finished,
            "speed": self.speed,
            "seed": sim.seed,
            "phase": self._phase_label(),
            "metrics": m,
            "agents": agents,
            "edges": self._edge_payload(),
            "nodes": self._node_payload(),
            "alerts": alerts,
            "bottlenecks": [b.as_dict() for b in bottlenecks],
            "primary_bottleneck": primary.as_dict() if primary else None,
            "prediction": {
                "source": self.predictor.source,
                "label": self.predictor.source_label,
                "horizons": list(self.settings.prediction.horizons_s),
                "top": self.predictor.summary(sim, limit=5),
            },
            "events": sim.event_log,
            "pending_events": self._pending_events(),
            "interventions": [
                {"strategy_id": a.strategy_id, "label": a.label, "t_s": a.t_s,
                 "agents_affected": a.agents_affected, "detail": a.detail}
                for a in sim.applied_interventions
            ],
            "reroute_paths": self._reroute_paths(),
            "error": self.last_error,
        }

    def _phase_label(self) -> str:
        t = self.sim.time
        label = self.scenario.phase_label
        for phase in self.venue_model.phases:
            end = phase.end_s if phase.end_s is not None else float("inf")
            if phase.start_s <= t < end:
                label = phase.name
        return label

    def _pending_events(self) -> list[dict[str, Any]]:
        out = []
        for i, ev in enumerate(self.scenario.timeline):
            if i in self.sim.fired_events:
                continue
            out.append({"index": i, "t_s": ev.t_s, "label": ev.label,
                        "detail": ev.detail, "severity": ev.severity,
                        "automatic": ev.automatic, "type": ev.type,
                        "target": ev.target, "factor": ev.factor})
        return out

    def _reroute_paths(self) -> list[dict[str, Any]]:
        """The alternative routes the crowd is actually being sent along.

        Only drawn once an intervention is live, and only for the diversion
        that matters: the paths leaving the congested corridor's upstream
        junction. Drawing every node whose adaptive hop happens to differ
        paints most of the venue green and tells the operator nothing.
        """
        if not self.sim.applied_interventions:
            return []
        primary = primary_bottleneck(self.sim, self.predictor.predict(self.sim))
        if primary is None:
            return []

        v = self.compiled
        edge_idx = primary.index
        decision_node = int(v.edge_src[edge_idx])
        upstream = {decision_node}
        for e in range(v.n_edges):
            if int(v.edge_dst[e]) == decision_node:
                upstream.add(int(v.edge_src[e]))

        out: list[dict[str, Any]] = []
        seen: set[tuple] = set()
        for slot, dest in enumerate(self.sim.dest_indices):
            for node_idx in sorted(upstream):
                base_hop = int(self.sim.tables.next_hop[POLICY_SHORTEST, slot, node_idx])
                adapt_hop = int(self.sim.tables.next_hop[POLICY_ADAPTIVE, slot, node_idx])
                if base_hop < 0 or adapt_hop < 0 or base_hop == adapt_hop:
                    continue
                _, edges = self.sim.tables.path_nodes(POLICY_ADAPTIVE, slot, node_idx)
                if not edges:
                    continue
                key = tuple(edges)
                if key in seen:
                    continue
                seen.add(key)
                out.append({
                    "from": v.node_ids[node_idx],
                    "to": v.node_ids[dest],
                    "edges": [v.edge_ids[e] for e in edges],
                    "base_edges": [v.edge_base_id[e] for e in edges],
                })
                if len(out) >= 3:
                    return out
        return out

    def summary(self) -> dict[str, Any]:
        return {
            "session_id": self.id,
            "kind": self.kind,
            "venue_id": self.config.venue_id,
            "scenario_id": self.config.scenario_id,
            "seed": self.sim.seed,
            "crowd_size": self.sim.n_agents,
            "speed": self.speed,
            "playing": self.playing,
            "finished": self.finished,
            "t_s": round(self.sim.time, 1),
            "duration_s": self.scenario.duration_s,
            "subscribers": self.broadcaster.count,
            "created_at": self.created_at,
            "config": self.config.as_dict(),
        }


class ReplaySession:
    """Plays back a precomputed run, exposing the same surface as a live one.

    This is the demo safety net. It is only used when a live session cannot be
    created, or when a recording is requested explicitly.
    """

    kind = "replay"

    def __init__(self, recording_path: Path, settings: Settings) -> None:
        self.id = uuid.uuid4().hex[:12]
        self.settings = settings
        self.created_at = time.time()
        with recording_path.open("r", encoding="utf-8") as fh:
            blob = json.load(fh)
        self.meta = blob["meta"]
        self.frames: list[dict[str, Any]] = blob["frames"]
        self.strategy_run: dict[str, Any] | None = blob.get("strategy_run")
        self.cursor = 0
        self.speed = int(self.meta.get("speed", 10))
        self.playing = False
        self.finished = False
        self.frame_index = 0
        self.last_error: str | None = None
        self.last_strategy_run = self.strategy_run
        self.broadcaster = Broadcaster()
        self.last_seen = time.time()
        self.venue_model = load_venue(self.meta["venue_id"])
        self.scenario = load_scenario(self.meta["scenario_id"])
        self._task: asyncio.Task | None = None
        self._applied: list[dict[str, Any]] = []

    def start_loop(self) -> None:
        if self._task is None or self._task.done():
            self._task = asyncio.create_task(self._run_loop())

    async def close(self) -> None:
        self.playing = False
        if self._task is not None:
            self._task.cancel()
            try:
                await self._task
            except (asyncio.CancelledError, Exception):
                pass

    async def _run_loop(self) -> None:
        interval = self.settings.server.frame_interval_s
        while True:
            if self.broadcaster.count == 0:
                await asyncio.sleep(0.4)
                continue
            self.last_seen = time.time()
            if self.playing and not self.finished:
                stride = max(1, int(round(self.speed / max(self.meta.get("speed", 10), 1))))
                self.cursor = min(self.cursor + stride, len(self.frames) - 1)
                if self.cursor >= len(self.frames) - 1:
                    self.finished = True
                    self.playing = False
                self.broadcaster.publish(self.frame())
            await asyncio.sleep(interval)

    def play(self) -> None:
        if not self.finished:
            self.playing = True

    def pause(self) -> None:
        self.playing = False

    def set_speed(self, speed: int) -> None:
        self.speed = int(min(max(speed, 1), max(SPEED_CHOICES)))

    async def step_once(self, seconds: float = 10.0) -> None:
        self.cursor = min(self.cursor + 1, len(self.frames) - 1)
        self.broadcaster.publish(self.frame())

    async def run_to(self, target_time_s: float) -> None:
        for i, f in enumerate(self.frames):
            if f["t_s"] >= target_time_s:
                self.cursor = i
                break
        else:
            self.cursor = len(self.frames) - 1
        self.broadcaster.publish(self.frame())

    def trigger_event(self, index: int) -> dict[str, Any]:
        return {"applied": False, "reason": "recorded run"}

    async def evaluate_strategies(self, horizon_s: float | None = None,
                                  strategy_ids: list[str] | None = None) -> dict[str, Any]:
        payload = self.strategy_run or {"available": False,
                                        "reason": "no recorded strategy run"}
        self.broadcaster.publish({"type": "strategy", "session_id": self.id,
                                  "payload": payload})
        return payload

    async def apply_strategy(self, strategy_id: str) -> dict[str, Any]:
        # Jump to the recorded post-intervention branch if one exists.
        branch = (self.meta.get("applied_branches") or {}).get(strategy_id)
        if branch is not None:
            self.cursor = min(int(branch), len(self.frames) - 1)
        self._applied.append({"strategy_id": strategy_id, "t_s": self.frames[self.cursor]["t_s"]})
        self.broadcaster.publish(self.frame())
        return {"applied": True, "strategy": {"id": strategy_id},
                "agents_affected": self.meta.get("agents_affected", 0),
                "t_s": self.frames[self.cursor]["t_s"]}

    def frame(self, include_agents: bool = True) -> dict[str, Any]:
        f = dict(self.frames[self.cursor])
        self.frame_index += 1
        f.update({"session_id": self.id, "kind": self.kind,
                  "frame": self.frame_index, "playing": self.playing,
                  "finished": self.finished, "speed": self.speed})
        if self._applied:
            f["interventions"] = self._applied
        return f

    def summary(self) -> dict[str, Any]:
        return {
            "session_id": self.id,
            "kind": self.kind,
            "venue_id": self.meta["venue_id"],
            "scenario_id": self.meta["scenario_id"],
            "seed": self.meta.get("seed"),
            "crowd_size": self.meta.get("crowd_size"),
            "speed": self.speed,
            "playing": self.playing,
            "finished": self.finished,
            "t_s": self.frames[self.cursor]["t_s"],
            "duration_s": self.scenario.duration_s,
            "subscribers": self.broadcaster.count,
            "created_at": self.created_at,
            "config": {"venue_id": self.meta["venue_id"],
                       "scenario_id": self.meta["scenario_id"],
                       "seed": self.meta.get("seed")},
        }


class SessionManager:
    """Creates, tracks and disposes of sessions."""

    def __init__(self, settings: Settings) -> None:
        self.settings = settings
        self.sessions: dict[str, SimulationSession | ReplaySession] = {}

    def get(self, session_id: str):
        return self.sessions.get(session_id)

    def list(self) -> list[dict[str, Any]]:
        return [s.summary() for s in self.sessions.values()]

    async def create(self, config: SessionConfig, allow_fallback: bool = True):
        await self.reap_idle()
        await self._evict_if_needed()
        try:
            session = SimulationSession(config, self.settings)
        except Exception as exc:
            if not (allow_fallback and self.settings.server.allow_fallback):
                raise
            recording = self._find_recording(config.scenario_id)
            if recording is None:
                raise
            session = ReplaySession(recording, self.settings)
            session.last_error = None
        self.sessions[session.id] = session
        session.start_loop()
        return session

    def create_replay(self, scenario_id: str) -> ReplaySession | None:
        recording = self._find_recording(scenario_id)
        if recording is None:
            return None
        session = ReplaySession(recording, self.settings)
        self.sessions[session.id] = session
        session.start_loop()
        return session

    def _find_recording(self, scenario_id: str) -> Path | None:
        path = FALLBACK_DIR / f"{scenario_id}.json"
        return path if path.exists() else None

    def has_recording(self, scenario_id: str) -> bool:
        return self._find_recording(scenario_id) is not None

    async def close(self, session_id: str) -> bool:
        session = self.sessions.pop(session_id, None)
        if session is None:
            return False
        await session.close()
        return True

    async def close_all(self) -> None:
        for sid in list(self.sessions):
            await self.close(sid)

    async def _evict_if_needed(self) -> None:
        limit = self.settings.server.max_sessions
        while len(self.sessions) >= limit:
            oldest = min(self.sessions.values(), key=lambda s: s.created_at)
            await self.close(oldest.id)

    async def reap_idle(self, grace_s: float = 90.0) -> int:
        """Dispose of sessions nobody has been watching for a while.

        A browser refresh abandons its session silently; without reaping, those
        accumulate for the length of the demo.
        """
        now = time.time()
        stale = [
            s.id for s in self.sessions.values()
            if s.broadcaster.count == 0 and (now - max(s.last_seen, s.created_at)) > grace_s
        ]
        for sid in stale:
            await self.close(sid)
        return len(stale)