"""BrowserSessionPool: the live browser sessions a single service replica owns.

Keyed by ``session_id`` (the id cloud-api carries in the conversation state), so
this replica is the holder for those sessions and dispatches their tool calls to
a per-session ``BrowserSessionThread``. Sessions are created lazily on first use
with a ``BrowserTools`` factory bound to the project, and evicted after an idle
TTL. One ``BrowserSessionPool`` per service process; the RabbitMQ transport feeds it.
"""

import threading
import time
from typing import Any, Callable, Dict, Optional, Tuple

from abstra_internals.cloud.browser_service.dispatch import dispatch
from abstra_internals.cloud.browser_service.session import BrowserSessionThread

# factory(project_id, editor_token) -> a BrowserTools instance for that project.
BrowserToolsProjectFactory = Callable[[str, Optional[str]], Any]

DEFAULT_IDLE_TTL_SECONDS = 300.0


class _Entry:
    __slots__ = ("thread", "project_id", "epoch", "last_used")

    def __init__(
        self,
        thread: BrowserSessionThread,
        project_id: str,
        epoch: int,
        last_used: float,
    ):
        self.thread = thread
        self.project_id = project_id
        # The highest epoch that has claimed this browser. A same-replica re-open
        # bumps it (reusing the browser, so tabs survive); a close from an OLDER
        # epoch is then a no-op, so a superseded holder can't destroy the browser
        # the current epoch is using.
        self.epoch = epoch
        self.last_used = last_used


class BrowserSessionPool:
    def __init__(
        self,
        browser_tools_factory: BrowserToolsProjectFactory,
        *,
        idle_ttl_seconds: float = DEFAULT_IDLE_TTL_SECONDS,
        clock: Callable[[], float] = time.monotonic,
        session_kwargs: Optional[Dict[str, Any]] = None,
    ) -> None:
        self._factory = browser_tools_factory
        self._ttl = idle_ttl_seconds
        self._clock = clock
        self._session_kwargs = session_kwargs or {}
        # Keyed by (project_id, session_id), NOT session_id alone. This is a
        # shared pool, so a command carrying (session_id=X, project_id=B) must
        # never bind to project A's live browser for the same session_id — a
        # different project is a different key, so it gets its own session.
        self._sessions: Dict[Tuple[str, str], _Entry] = {}
        self._lock = threading.Lock()

    def _get_or_create(
        self, session_id: str, project_id: str, editor_token: Optional[str], epoch: int
    ) -> _Entry:
        key = (project_id, session_id)
        with self._lock:
            entry = self._sessions.get(key)
            if entry is None:
                thread = BrowserSessionThread(
                    lambda: self._factory(project_id, editor_token),
                    **self._session_kwargs,
                )
                entry = _Entry(thread, project_id, epoch, self._clock())
                self._sessions[key] = entry
            elif epoch > entry.epoch:
                # A re-open landed on this same replica: keep the browser (its
                # tabs survive the re-open) but hand ownership to the new epoch.
                entry.epoch = epoch
            entry.last_used = self._clock()
            return entry

    def handle(
        self,
        session_id: str,
        project_id: str,
        editor_token: Optional[str],
        tool_name: str,
        args: Optional[Dict[str, Any]],
        epoch: int = 1,
    ) -> str:
        """Dispatch one browser tool call for a session, returning the
        pod-identical serialized string. The (long) browser call runs outside
        the registry lock, so sessions dispatch concurrently."""
        entry = self._get_or_create(session_id, project_id, editor_token, epoch)
        return dispatch(entry.thread, tool_name, args)

    def close(self, session_id: str, project_id: str, epoch: int = 1) -> None:
        """Close a session's browser, but ONLY if this epoch still owns it. A
        superseded holder (older epoch) calling close after a same-replica
        re-open must not tear down the browser the newer epoch is using."""
        key = (project_id, session_id)
        with self._lock:
            entry = self._sessions.get(key)
            if entry is None or entry.epoch != epoch:
                return  # stale/foreign close — the current epoch keeps the browser
            del self._sessions[key]
        entry.thread.close()

    def evict_idle(self) -> int:
        """Close sessions idle longer than the TTL. Returns how many were
        evicted. Called periodically by the transport's sweeper."""
        now = self._clock()
        with self._lock:
            expired_keys = [
                key
                for key, e in self._sessions.items()
                if now - e.last_used > self._ttl
            ]
            expired = [self._sessions.pop(key) for key in expired_keys]
        for e in expired:
            e.thread.close()
        return len(expired)

    def close_all(self) -> None:
        with self._lock:
            entries = list(self._sessions.values())
            self._sessions.clear()
        for e in entries:
            e.thread.close()

    @property
    def active_sessions(self) -> int:
        with self._lock:
            return len(self._sessions)
