"""Process-global registry of active WebSocket sessions.

Lets the HTTP control server (``server.py``) look up a running session
by id and inject commands into its recv loop without going through the
WebSocket. The MCP server uses this to drive the demo while the
browser's own WebSocket connection stays primary.

Torch-free: imported by ``server.py`` regardless of ``--no-backend``.
"""

from __future__ import annotations

import secrets
import threading
import time
from dataclasses import dataclass
from typing import Any, Callable, Optional


InjectFn = Callable[[dict, Optional[bytes]], None]
"""Push a control message into a session.

Args:
    data: parsed JSON message in the same shape the browser would send
        over the WebSocket (e.g. ``{"type": "prompt", "tags": "..."}``).
    audio: optional binary frame (channels/samples header + interleaved
        float32 PCM) for messages that need it (``set_timbre_source``,
        ``set_structure_source``, ``swap_source``).
"""


SnapshotFn = Callable[[], dict]
"""Return a JSON-serialisable snapshot of the session's current state."""


@dataclass
class SessionHandle:
    id: str
    started_at: float
    inject: InjectFn
    snapshot: SnapshotFn


_sessions: dict[str, SessionHandle] = {}
_lock = threading.Lock()


def new_session_id() -> str:
    """8-char URL-safe id. Generated by the backend before registering."""
    return secrets.token_urlsafe(6)


def register(handle: SessionHandle) -> None:
    with _lock:
        _sessions[handle.id] = handle


def unregister(session_id: str) -> None:
    with _lock:
        _sessions.pop(session_id, None)


def get(session_id: str) -> Optional[SessionHandle]:
    with _lock:
        return _sessions.get(session_id)


def list_handles() -> list[SessionHandle]:
    """Snapshot of all live sessions, newest first."""
    with _lock:
        return sorted(_sessions.values(), key=lambda h: h.started_at, reverse=True)


def most_recent() -> Optional[SessionHandle]:
    handles = list_handles()
    return handles[0] if handles else None
