"""Per-session event sequencing + bounded replay for WS reconnects.

Every event frame through :func:`server.write_json` (hence ``_emit``) gets a per-session monotonic
``seq`` and lands in a small ring per session; a reconnecting client calls ``session.events.since``
with its last seen seq and gets everything newer. Invariants: stdio TUI unaffected (``seq`` only on
event frames; Ink ignores unknown keys); one lock guards counters + buffers, and write_json already
serializes per-transport writes so stamping cannot reorder frames; memory bound =
_REPLAY_BUFFER_MAX events AND _REPLAY_BUFFER_BYTES_MAX frozen UTF-8 payload bytes per session,
_REPLAY_PROCESS_BYTES_MAX frozen UTF-8 payload bytes across at most _REPLAY_SESSIONS_MAX sessions,
oldest evicted FIFO (their seq/truncation counters are retained so a revisited session stays
monotonic within the epoch, #100122). Evicted or never-retained (oversized) frames leave a truncation watermark so a
reconnecting client refetches instead of trusting a replay with holes. A torn-down session's ring is freed at
teardown (``forget_session``) and never rebuilt by late frames.
"""

from __future__ import annotations

import json
import threading
import uuid
from collections import OrderedDict, deque

# Seq counters live in-process, so a restart resets them to 1 while clients hold high
# watermarks — events_since(sid, 97) would return [] with truncated=False forever. The
# epoch lets clients detect the restart and reset their watermarks.
_REPLAY_EPOCH = uuid.uuid4().hex

# A long turn emits ~hundreds of token events; 512 covers minutes of streaming plus
# all control events. Desktop users rarely exceed a dozen live chats.
_REPLAY_BUFFER_MAX = 512
_REPLAY_SESSIONS_MAX = 64
# A ring may legitimately hold many bounded 64 KiB tool results (512 of them ≈ 32 MiB per
# session, ×64 sessions before any cap); bound the frozen serialized payload so replay memory
# cannot scale with mutable Python object graphs after admission.
_REPLAY_BUFFER_BYTES_MAX = 4 * 1024 * 1024
_REPLAY_PROCESS_BYTES_MAX = 64 * 1024 * 1024

_replay_lock = threading.Lock()
# sid -> deque of (seq, frozen params JSON bytes without the server-owned seq, retained bytes).
# Bytes detach the replay ring from the live event object: later caller mutation cannot grow or
# rewrite an already-accounted cache entry.
_replay_buffers: OrderedDict[str, deque] = OrderedDict()
_replay_buffer_bytes: dict[str, int] = {}
_replay_evicted_through: dict[str, int] = {}
_replay_total_bytes = 0
_replay_next_seq: dict[str, int] = {}
# Torn-down sids (forget_session). A frame stamped later (a turn thread that outlived the teardown join, a
# compute-host bridge echo) still gets a seq and a watermark but is never retained, so it cannot rebuild a
# ring that nothing would free until FIFO eviction. One str per closed session, like _replay_next_seq.
_replay_forgotten: set[str] = set()


def replay_epoch() -> str:
    """Opaque token identifying this server process's seq numbering."""
    return _REPLAY_EPOCH


def _freeze_params(params: dict) -> bytes:
    """Serialize one replay snapshot before the caller can mutate its object graph."""
    return json.dumps(params, ensure_ascii=False, separators=(",", ":")).encode(
        "utf-8", errors="surrogatepass"
    )


def _thaw_params(payload: bytes, seq: int) -> dict:
    """Rebuild a client event object and restore the server-owned sequence number."""
    event = json.loads(payload.decode("utf-8", errors="surrogatepass"))
    event["seq"] = seq
    return event


def _stamp_event(obj: dict) -> None:
    """Stamp one outgoing event frame (mutates obj in place) and record a frozen snapshot."""
    if obj.get("method") != "event":
        return
    params = obj.get("params")
    if not isinstance(params, dict):
        return
    sid = params.get("session_id") or ""
    if not sid:
        # Session-less global events (skin.changed etc.) are re-fetchable via their own RPCs.
        return
    # Freeze OUTSIDE the lock (same rule as transport.write) so one large payload cannot stall
    # other threads' frames. The server-owned seq is restored when replaying, so the snapshot can
    # be serialized before sequence assignment without retaining the mutable params dict.
    payload = _freeze_params(params)
    size = len(payload)
    with _replay_lock:
        global _replay_total_bytes
        seq = _replay_next_seq.get(sid, 0) + 1
        _replay_next_seq[sid] = seq
        params["seq"] = seq
        if sid in _replay_forgotten:
            _replay_evicted_through[sid] = seq
            return
        buf = _replay_buffers.get(sid)
        if buf is None:
            buf = _replay_buffers[sid] = deque()
            _replay_buffer_bytes[sid] = 0
            while len(_replay_buffers) > _REPLAY_SESSIONS_MAX:
                _drop_ring_locked(next(iter(_replay_buffers)))
        if size > _REPLAY_BUFFER_BYTES_MAX or size > _REPLAY_PROCESS_BYTES_MAX:
            _replay_evicted_through[sid] = seq
            return
        buf.append((seq, payload, size))
        _replay_buffer_bytes[sid] += size
        _replay_total_bytes += size
        while len(buf) > _REPLAY_BUFFER_MAX or _replay_buffer_bytes[sid] > _REPLAY_BUFFER_BYTES_MAX:
            evicted_seq, _event, evicted_size = buf.popleft()
            _replay_buffer_bytes[sid] -= evicted_size
            _replay_total_bytes -= evicted_size
            _replay_evicted_through[sid] = max(_replay_evicted_through.get(sid, 0), evicted_seq)
        while _replay_total_bytes > _REPLAY_PROCESS_BYTES_MAX:
            for evict_sid, evict_buf in _replay_buffers.items():
                if evict_buf:
                    evicted_seq, _event, evicted_size = evict_buf.popleft()
                    _replay_buffer_bytes[evict_sid] -= evicted_size
                    _replay_total_bytes -= evicted_size
                    _replay_evicted_through[evict_sid] = max(
                        _replay_evicted_through.get(evict_sid, 0), evicted_seq
                    )
                    break


def _drop_ring_locked(sid: str) -> None:
    """Free one session's retained frames (FIFO cap eviction and teardown). Caller holds ``_replay_lock``.

    The ring goes, the session's seq numbering stays (#100122): a revisited session continues from its high
    seq instead of restarting at 1 under clients' still-held watermarks (a reset seq is invisible to
    dispatchIfNewer, so replay AND live frames silently vanish). The truncation watermark is raised to the
    latest stamped seq: a client holding any older watermark has a gap and must refetch history instead of
    trusting the new tail. Both counters are one int per session id seen this process."""
    global _replay_total_bytes
    if _replay_buffers.pop(sid, None) is None:
        return
    _replay_total_bytes -= _replay_buffer_bytes.pop(sid, 0)
    _replay_evicted_through[sid] = max(_replay_evicted_through.get(sid, 0), _replay_next_seq.get(sid, 0))


def forget_session(sid: str) -> None:
    """Release a torn-down session's replay ring. A runtime sid is never reattached after teardown
    (resume mints a new one), so its frames can only be replayed to nobody."""
    if not sid:
        return
    with _replay_lock:
        _replay_forgotten.add(sid)
        _drop_ring_locked(sid)


def events_since(sid: str, last_seen: int) -> list[dict]:
    """Recorded EVENT OBJECTS (each frame's ``params`` dict) with seq > last_seen for *sid*.

    Returning the full JSON-RPC envelope would make every replayed event fail the
    client's ``event.type`` gate and be silently dropped. Decode after releasing the ring lock so a
    large replay cannot block concurrent event stamping.
    """
    with _replay_lock:
        buf = _replay_buffers.get(sid or "")
        snapshots = [(seq, payload) for seq, payload, _size in buf if seq > last_seen] if buf else []
    return [_thaw_params(payload, seq) for seq, payload in snapshots]


def is_truncated(sid: str, last_seen: int) -> bool:
    """True when events between *last_seen* and the ring's oldest retained seq were
    evicted — the client must refetch history instead of trusting the replay."""
    with _replay_lock:
        return last_seen < _replay_evicted_through.get(sid or "", 0)


def latest_seq(sid: str) -> int:
    """Current highest stamped seq for *sid* (0 when unknown)."""
    with _replay_lock:
        return _replay_next_seq.get(sid or "", 0)


def reset_replay_state() -> None:
    """Test hook."""
    with _replay_lock:
        global _replay_total_bytes
        _replay_buffers.clear()
        _replay_buffer_bytes.clear()
        _replay_evicted_through.clear()
        _replay_next_seq.clear()
        _replay_forgotten.clear()
        _replay_total_bytes = 0


def replay_stats() -> dict:
    """Telemetry: buffer occupancy for the ops/debug surface."""
    with _replay_lock:
        return {
            "sessions": len(_replay_buffers),
            "events": sum(len(buffer) for buffer in _replay_buffers.values()),
            "bytes": _replay_total_bytes,
            "max_per_session": _REPLAY_BUFFER_MAX,
            "max_bytes_per_session": _REPLAY_BUFFER_BYTES_MAX,
            "max_bytes_process": _REPLAY_PROCESS_BYTES_MAX,
        }