"""Global session registry for managing FriskSession instances by UUID.

This module provides a thread-safe registry for storing and retrieving
FriskSession objects using UUID identifiers. Sessions are automatically
registered when created and can be unregistered when ended.
"""

from __future__ import annotations
import threading
from typing import TYPE_CHECKING, Any
from uuid import UUID

if TYPE_CHECKING:
    from .frisk_session import FriskSession


class SessionNotFoundError(Exception):
    """Raised when a session UUID is not found in the registry.

    This can occur when:
    - The UUID is invalid or not registered
    - The session has already ended and been cleaned up
    - The session was created in a different runtime instance
    """

    def __init__(self, session_id: UUID | str):
        self.session_id = session_id
        message = f"""\
Session not found: {session_id}

The requested session UUID does not exist in the registry.
This can happen if:
  - The session has already ended (via end_tracing())
  - The UUID is invalid or was never registered
  - The session was created in a different process/runtime

Please ensure:
  1. You're using a valid session ID returned from Frisk.session()
  2. The session has not been ended yet
  3. You're in the same runtime where the session was created
"""
        super().__init__(message)


class SessionRegistry:
    """Thread-safe global registry for FriskSession instances.

    This singleton registry manages the lifecycle of sessions, allowing them
    to be stored and retrieved by UUID across the application without passing
    FriskSession objects directly.
    """

    _instance = None
    _lock = threading.Lock()
    _sessions: dict[str, "FriskSession[Any]"]
    _sessions_lock: threading.Lock

    def __new__(cls):
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super().__new__(cls)
                    cls._instance._sessions = {}
                    cls._instance._sessions_lock = threading.Lock()
        return cls._instance

    def register(self, session_id: str, session: FriskSession[Any]) -> None:
        """Register a session in the registry.

        Args:
            session_id: UUID identifier for the session
            session: FriskSession instance to register
        """
        with self._sessions_lock:
            self._sessions[session_id] = session

    def get(self, session_id: str) -> FriskSession[Any]:
        """Retrieve a session from the registry.

        Args:
            session_id: UUID or string representation of the session ID

        Returns:
            The FriskSession instance

        Raises:
            SessionNotFoundError: If the session is not found in the registry
        """
        with self._sessions_lock:
            session = self._sessions.get(session_id)
            if session is None:
                raise SessionNotFoundError(session_id)
            return session

    def unregister(self, session_id: str) -> None:
        """Remove a session from the registry.

        Args:
            session_id: UUID identifier of the session to remove
        """
        with self._sessions_lock:
            self._sessions.pop(session_id, None)

    def clear(self) -> None:
        """Clear all sessions from the registry.

        This is primarily useful for testing.
        """
        with self._sessions_lock:
            self._sessions.clear()
