"""Transport for the browser service: RabbitMQ opens + NATS session channels.

cloud-api and the service live in different clusters with no direct route; both
reach the shared RabbitMQ broker AND the shared NATS server. The two legs use
different rails on purpose:

  1. Opens ride a durable RabbitMQ queue: cloud-api publishes
     {session_id, project_id, epoch}; replicas are competing consumers, so
     exactly one becomes the holder — and a FULL replica can nack-requeue
     (the admission backpressure). Queue semantics are load-bearing here.
  2. The session channel rides NATS pub/sub subjects (epoch-scoped): the holder
     subscribes to the cmd subject, publishes a READY reply, then serves
     {corr_id, tool_name, args} -> {corr_id, ok, result}. ``__close__`` ends
     the session. Pub/sub is deliberate: per-session RabbitMQ queues/channels
     are exactly the broker-object cardinality that saturated the broker's
     mgmt CPU before (the forms lesson) — NATS subjects have no lifecycle, no
     debris, no per-session broker objects.

Load-bearing contracts:

  - READY handshake: cloud-api subscribes its reply subject, THEN publishes the
    open, and sends NO command until the holder's ready reply — pub/sub drops
    anything published before the subscriber exists, and under backpressure the
    holder may attach minutes later.
  - ``__close__`` from cloud-api is the PRIMARY reclamation path; the idle
    timeout is only a crash backstop — a missed close pins a Chromium that long.
  - The ``epoch`` is the fencing token: it scopes the subject names, so a
    re-open (client command timeout) gets DISTINCT subjects and a stale-but-
    alive holder can never split a session's commands with the new one.
  - ``max_sessions`` + prefetch(1) is the admission control: a full replica
    requeues opens (short backoff), leaving them as backpressure for a replica
    with room — the deployment's memory HPA grows the fleet on the same signal.

Still owed before heavy traffic: adopt ``RabbitMQConsumer`` in ``run()`` —
``stop()`` is a bare flag (no reconnect/heartbeat, SIGTERM can hang consuming).
"""

import json
import re
import threading
import time
from dataclasses import dataclass
from typing import Any, Dict, Optional, Protocol

from abstra_internals.logger import AbstraLogger

# How long a full replica pauses after requeueing an open, so two full replicas
# don't ping-pong the message at broker speed while the autoscaler reacts.
FULL_REQUEUE_BACKOFF_SECONDS = 0.5

# Shared durable queue cloud-api publishes session-open requests to.
BROWSER_SERVICE_OPEN_QUEUE = "browser-service:open"

# Reserved corr_id of the holder's ready reply (the handshake).
READY_CORR_ID = "__ready__"


def _subject_token(session_id: str) -> str:
    """Session ids are NATS subject tokens; strip anything that isn't."""
    return re.sub(r"[^A-Za-z0-9_-]", "_", session_id)


# Per-session subject names, epoch-scoped (cloud-api mirrors these).
def cmd_subject(session_id: str, epoch: int) -> str:
    return f"browser-service.cmd.{_subject_token(session_id)}.{epoch}"


def reply_subject(session_id: str, epoch: int) -> str:
    return f"browser-service.reply.{_subject_token(session_id)}.{epoch}"


# Sentinel tool name cloud-api sends to end a session.
CLOSE_TOOL = "__close__"


@dataclass(frozen=True)
class SessionOpen:
    session_id: str
    project_id: str
    # Fencing token: cloud-api assigns 1 at first open, bumps on re-open.
    epoch: int = 1
    # Deliberately no credentials: this message rides a durable queue.
    # (open_page's editor_auth token rides its bounce URL on the ephemeral
    # cmd queue instead.)

    @staticmethod
    def from_dict(d: Dict[str, Any]) -> "SessionOpen":
        return SessionOpen(
            session_id=str(d["session_id"]),
            project_id=str(d["project_id"]),
            epoch=int(d.get("epoch", 1)),
        )


class SessionPoolLike(Protocol):
    def handle(
        self,
        session_id: str,
        project_id: str,
        editor_token: Optional[str],
        tool_name: str,
        args: Optional[Dict[str, Any]],
        epoch: int,
    ) -> str: ...

    def close(self, session_id: str, project_id: str, epoch: int) -> None: ...


def handle_command(
    pool: SessionPoolLike, session: SessionOpen, command: Dict[str, Any]
) -> Dict[str, Any]:
    """One command message -> one reply message (pure). Tool errors become the
    tunneled-convention result text, so cloud-api relays ``result`` verbatim."""
    corr_id = command.get("corr_id")
    tool_name = command.get("tool_name", "")
    if tool_name == CLOSE_TOOL:
        return {"corr_id": corr_id, "ok": True, "closed": True, "result": ""}
    args = command.get("args") or {}
    try:
        result = pool.handle(
            session.session_id,
            session.project_id,
            None,  # editor_token is deferred to Phase 4 (browser_open_page)
            tool_name,
            args,
            session.epoch,
        )
        return {"corr_id": corr_id, "ok": True, "result": result}
    except Exception as e:  # noqa: BLE001 - surfaced to the model as tool output
        return {"corr_id": corr_id, "ok": False, "result": f"Error executing tool: {e}"}


class _NatsSessionChannel:
    """Session channel over NATS subjects, presenting the same blocking
    recv()/send()/close() protocol ``_serve_session`` consumes. recv() returns
    None on idle timeout or a closed connection — what queue expiry used to
    signal — so an abandoned session's serve thread ends instead of leaking."""

    def __init__(self, inner: Any, idle_seconds: float) -> None:
        self._inner = inner
        self._idle_seconds = idle_seconds

    def recv(self) -> Any:
        try:
            return self._inner.recv(timeout=self._idle_seconds)
        except (TimeoutError, EOFError):
            return None

    def send(self, obj: Any) -> None:
        self._inner.send(obj)

    def close(self) -> None:
        self._inner.close()


class BrowserServer:
    """Consumes session-open requests (RabbitMQ) and serves each session's
    command channel (NATS subjects)."""

    def __init__(
        self,
        connection_uri: str,
        pool: SessionPoolLike,
        max_sessions: Optional[int] = None,
        nats_url: Optional[str] = None,
        nats_creds: Optional[str] = None,
        session_idle_seconds: float = 7_500.0,
    ) -> None:
        self._uri = connection_uri
        self._pool = pool
        self._stopped = False
        self._max_sessions = max_sessions  # None = uncapped (tests)
        self._active_sessions = 0
        self._active_lock = threading.Lock()
        self._nats_url = nats_url
        self._nats_creds = nats_creds
        self._session_idle_seconds = session_idle_seconds
        self._nats_persistent: Any = None
        self._nats_lock = threading.Lock()

    @property
    def active_sessions(self) -> int:
        with self._active_lock:
            return self._active_sessions

    def try_take_session(self, session: SessionOpen) -> bool:
        """Become the session's holder if this replica has room; False = caller
        requeues the open. A channel-open failure also returns False (requeue)
        rather than killing the open consumer on one broker hiccup."""
        with self._active_lock:
            if (
                self._max_sessions is not None
                and self._active_sessions >= self._max_sessions
            ):
                return False
            self._active_sessions += 1
        try:
            conn = self.open_session_channel(session)
            # Handshake: cloud-api sends no command until this arrives (pub/sub
            # drops anything published before our subscription existed).
            conn.send({"corr_id": READY_CORR_ID, "ok": True, "result": ""})
        except Exception as e:  # noqa: BLE001
            with self._active_lock:
                self._active_sessions = max(0, self._active_sessions - 1)
            AbstraLogger.error(
                f"[browser-service] could not open session channel for "
                f"{session.session_id}:{session.epoch}: {e}"
            )
            return False
        threading.Thread(
            target=self._serve_session,
            args=(conn, session),
            name=f"browser-session-{session.session_id[:8]}",
            daemon=True,
        ).start()
        return True

    def _serve_session(self, conn: Any, session: SessionOpen) -> None:
        """Loop over a session's command channel until closed."""
        try:
            while not self._stopped:
                command = conn.recv()
                if command is None:
                    break
                if isinstance(command, (str, bytes)):
                    command = json.loads(command)
                reply = handle_command(self._pool, session, command)
                conn.send(reply)
                if reply.get("closed"):
                    break
        except Exception as e:  # noqa: BLE001
            AbstraLogger.warning(
                f"[browser-service] session {session.session_id} loop ended: {e}"
            )
        finally:
            # Epoch-guarded: if a newer epoch re-opened on this same replica and
            # took over the browser, this (older) close is a no-op.
            self._pool.close(session.session_id, session.project_id, session.epoch)
            try:
                conn.close()
            except Exception:  # noqa: BLE001
                pass
            with self._active_lock:
                self._active_sessions = max(0, self._active_sessions - 1)

    def _get_nats_persistent(self) -> Any:
        """One NATS TCP connection per replica; per-session subscriptions ride it."""
        from abstra_internals.utils.nats_connection import NATSPersistentConnection

        with self._nats_lock:
            if self._nats_persistent is None or not self._nats_persistent.is_alive:
                if not self._nats_url or not self._nats_creds:
                    raise RuntimeError(
                        "[browser-service] NATS url/creds not configured; the "
                        "session channel requires ABSTRA_NATS_URL + ABSTRA_NATS_CREDS"
                    )
                self._nats_persistent = NATSPersistentConnection(
                    self._nats_url, self._nats_creds
                )
            return self._nats_persistent

    def open_session_channel(self, session: SessionOpen) -> Any:
        """Holder-side channel: RECV cmd subject, SEND reply subject. Split out
        so tests can substitute a fake."""
        from abstra_internals.utils.nats_connection import NATSConnection

        inner = NATSConnection(
            nats_url=self._nats_url or "",
            nats_creds=self._nats_creds or "",
            send_subject=reply_subject(session.session_id, session.epoch),
            recv_subject=cmd_subject(session.session_id, session.epoch),
            execution_id=f"{session.session_id}:{session.epoch}",
            persistent=self._get_nats_persistent(),
        )
        return _NatsSessionChannel(inner, self._session_idle_seconds)

    def run(self) -> None:
        """Consume the shared open queue; spawn a serving thread per session."""
        import pika

        AbstraLogger.warning(
            f"[browser-service] server starting; consuming open queue "
            f"(max sessions: {self._max_sessions or 'uncapped'})"
        )
        connection = pika.BlockingConnection(pika.URLParameters(self._uri))
        channel = connection.channel()
        channel.queue_declare(queue=BROWSER_SERVICE_OPEN_QUEUE, durable=True)
        # Never buffer opens this replica may not have capacity to serve.
        channel.basic_qos(prefetch_count=1)

        def on_open(ch: Any, method: Any, _props: Any, body: bytes) -> None:
            try:
                session = SessionOpen.from_dict(json.loads(body))
            except Exception as e:  # noqa: BLE001 - drop a malformed open, keep serving
                AbstraLogger.error(f"[browser-service] bad open message dropped: {e}")
                ch.basic_ack(method.delivery_tag)
                return
            if self.try_take_session(session):
                ch.basic_ack(method.delivery_tag)
                return
            # Full (or channel-open hiccup): back to the queue for a replica
            # with room. Queued opens are the fleet's backpressure.
            AbstraLogger.warning(
                f"[browser-service] at capacity ({self.active_sessions}); "
                f"requeueing open for {session.session_id}:{session.epoch}"
            )
            ch.basic_nack(method.delivery_tag, requeue=True)
            time.sleep(FULL_REQUEUE_BACKOFF_SECONDS)

        channel.basic_consume(BROWSER_SERVICE_OPEN_QUEUE, on_open)
        try:
            channel.start_consuming()
        finally:
            with_suppress = getattr(connection, "close", lambda: None)
            try:
                with_suppress()
            except Exception:  # noqa: BLE001
                pass

    def stop(self) -> None:
        self._stopped = True
        with self._nats_lock:
            if self._nats_persistent is not None:
                try:
                    self._nats_persistent.close()
                except Exception:  # noqa: BLE001
                    pass
                self._nats_persistent = None
