from threading import Lock
from typing import Any

from flask import Flask, request, session
from flask_socketio import SocketIO, join_room
from loguru import logger

from ...constants import ResearchStatus
from ...database.encrypted_db import db_manager
from ...database.session_passwords import session_password_store
from ..routes.globals import get_active_research_snapshot


def _install_origin_rejection_logging(socketio: SocketIO) -> bool:
    """Re-emit engine.io's silenced WebSocket origin rejections via loguru.

    engine.io validates the Origin at handshake and calls
    ``_log_error_once('<origin> is not an accepted origin.', 'bad-origin')``,
    but the server runs with ``logger=False`` so that message never surfaces —
    the only symptom of a misconfigured WebSocket origin is a frozen progress
    UI. Wrap that one call to log a WARNING (deduped per origin) pointing at the
    fix. An Origin is a scheme+host, not PII. Best-effort: a no-op (returns
    False) if engine.io internals change, so it can never break startup.

    The dedup set is capped: the handshake is pre-auth and ``Origin`` is
    attacker-controlled, so an unbounded set would be a memory-growth + log-
    amplification vector. After ``cap`` distinct origins we stop tracking/warning
    (an operator has more than enough signal by then).
    """
    try:
        eio = socketio.server.eio
        original = eio._log_error_once
    except AttributeError:
        logger.debug(
            "Socket.IO: origin-rejection logging not installed "
            "(engine.io internals changed); handshake rejections stay silent"
        )
        return False

    warned: set[str] = set()
    cap = 100

    def _log_error_once(message, message_key):
        if (
            message_key == "bad-origin"
            and len(warned) < cap
            and message not in warned
        ):
            warned.add(message)
            logger.warning(
                f"Socket.IO rejected a WebSocket handshake: {message} Set "
                "LDR_SECURITY_WEBSOCKET_ALLOWED_ORIGINS to this origin if it is "
                "your front-end; behind a TLS-terminating proxy, also forward "
                "X-Forwarded-Proto so the same-origin check sees https."
            )
        return original(message, message_key)

    eio._log_error_once = _log_error_once
    return True


class SocketIOService:
    """
    Singleton class for managing SocketIO connections and subscriptions.
    """

    _instance = None

    def __new__(cls, *args: Any, app: Flask | None = None, **kwargs: Any):
        """
        Args:
            app: The Flask app to bind this service to. It must be specified
                the first time this is called and the singleton instance is
                created, but will be ignored after that.
            *args: Arguments to pass to the superclass's __new__ method.
            **kwargs: Keyword arguments to pass to the superclass's __new__ method.
        """
        if not cls._instance:
            if app is None:
                raise ValueError(
                    "Flask app must be specified to create a SocketIOService instance."
                )
            cls._instance = super(SocketIOService, cls).__new__(
                cls, *args, **kwargs
            )
            cls._instance.__init_singleton(app)
        return cls._instance

    def __init_singleton(self, app: Flask) -> None:
        """
        Initializes the singleton instance.

        Args:
            app: The app to bind this service to.

        """
        self.__app = app  # Store the Flask app reference

        # Determine WebSocket CORS policy from env var or default
        from ...settings.env_registry import get_env_setting

        ws_origins_env = get_env_setting("security.websocket.allowed_origins")
        socketio_cors: str | list[str] | None
        if ws_origins_env is not None:
            if ws_origins_env == "*":
                socketio_cors = "*"
            elif ws_origins_env:
                socketio_cors = [o.strip() for o in ws_origins_env.split(",")]
            else:
                socketio_cors = None
        else:
            # No env var set — fail closed to same-origin only, matching HTTP CORS default
            socketio_cors = None

        if socketio_cors is None:
            logger.info(
                "Socket.IO CORS: same-origin only (set LDR_SECURITY_WEBSOCKET_ALLOWED_ORIGINS to configure)"
            )
        elif socketio_cors == "*":
            logger.debug("Socket.IO CORS: all origins allowed")
        else:
            logger.info(f"Socket.IO CORS: restricted to {socketio_cors}")

        self.__socketio = SocketIO(
            app,
            cors_allowed_origins=socketio_cors,
            async_mode="threading",
            path="/socket.io",
            logger=False,
            engineio_logger=False,
            ping_timeout=20,
            ping_interval=5,
        )

        # Make a rejected WebSocket origin diagnosable (otherwise it is a silent
        # frozen progress UI). Skipped for the allow-all case, which rejects
        # nothing.
        if socketio_cors != "*":
            _install_origin_rejection_logging(self.__socketio)

        # Socket subscription tracking.
        self.__socket_subscriptions: dict[str, Any] = {}
        # Set to false to disable logging in the event handlers. This can
        # be necessary because it will sometimes run the handlers directly
        # during a call to `emit` that was made in a logging handler.
        self.__logging_enabled = True
        # Protects access to shared state.
        self.__lock = Lock()

        # Register events.
        @self.__socketio.on("connect")
        def on_connect():
            return self.__handle_connect(request)

        @self.__socketio.on("disconnect")
        def on_disconnect(reason: str):
            self.__handle_disconnect(request, reason)

        @self.__socketio.on("subscribe_to_research")
        def on_subscribe(data):
            self.__handle_subscribe(data, request)

        # Backwards-compatible alias: the JS client emits 'join' on subscribe.
        # Without this, the catch-up snapshot in __handle_subscribe never
        # fires and per-client targeting falls through to broadcast.
        @self.__socketio.on("join")
        def on_join(data):
            self.__handle_subscribe(data, request)

        @self.__socketio.on("leave")
        def on_leave(data):
            self.__handle_unsubscribe(data, request)

        @self.__socketio.on("unsubscribe_from_research")
        def on_unsubscribe(data):
            self.__handle_unsubscribe(data, request)

        @self.__socketio.on_error
        def on_error(e):
            return self.__handle_socket_error(e)

        @self.__socketio.on_error_default
        def on_default_error(e):
            return self.__handle_default_error(e)

    def __log_info(self, message: str, *args: Any, **kwargs: Any) -> None:
        """Log an info message."""
        if self.__logging_enabled:
            logger.info(message, *args, **kwargs)

    def __log_error(self, message: str, *args: Any, **kwargs: Any) -> None:
        """Log an error message."""
        if self.__logging_enabled:
            logger.error(message, *args, **kwargs)

    def __log_exception(self, message: str, *args: Any, **kwargs: Any) -> None:
        """Log an exception."""
        if self.__logging_enabled:
            logger.exception(message, *args, **kwargs)

    @staticmethod
    def user_room(username: str) -> str:
        """Socket.IO room name that every one of a user's connected tabs joins.

        Used to scope user-private events to a single account. Kept here so the
        connect handler and event emitters share one definition and cannot
        drift apart.
        """
        return f"user:{username}"

    def emit_socket_event(self, event, data, room=None):
        """
        Emit a socket event to clients.

        Args:
            event: The event name to emit
            data: The data to send with the event
            room: Optional room ID to send to specific client

        Returns:
            bool: True if emission was successful, False otherwise
        """
        try:
            # If room is specified, only emit to that room
            if room:
                self.__socketio.emit(event, data, room=room)
            else:
                # Otherwise broadcast to all
                self.__socketio.emit(event, data)
            return True
        except Exception:
            logger.exception(f"Error emitting socket event {event}")
            return False

    def emit_to_subscribers(
        self, event_base, research_id, data, enable_logging: bool = True
    ):
        """
        Emit an event to all subscribers of a specific research.

        Args:
            event_base: Base event name (will be formatted with research_id)
            research_id: ID of the research
            data: The data to send with the event
            enable_logging: If set to false, this will disable all logging,
                which is useful if we are calling this inside of a logging
                handler.

        Returns:
            bool: True if emission was successful, False otherwise

        """
        if not enable_logging:
            self.__logging_enabled = False

        try:
            full_event = f"{event_base}_{research_id}"

            # Emit only to specific subscribers (no broadcast) to avoid
            # duplicate messages and reduce server load under concurrency
            with self.__lock:
                subscriptions = self.__socket_subscriptions.get(research_id)
                if subscriptions:
                    subscriptions = (
                        subscriptions.copy()
                    )  # snapshot avoids RuntimeError
                else:
                    subscriptions = None
            if subscriptions is not None:
                for sid in subscriptions:
                    try:
                        self.__socketio.emit(full_event, data, room=sid)
                    except Exception:
                        self.__log_exception(
                            f"Error emitting to subscriber {sid}"
                        )
            # When no targeted subscribers exist yet, drop the event.
            # The catch-up snapshot in __handle_subscribe replays the
            # latest progress on subscribe, so early-arriving events
            # are recovered correctly without a cross-user broadcast.

            return True
        except Exception:
            self.__log_exception(
                f"Error emitting to subscribers for research {research_id}"
            )
            return False
        finally:
            self.__logging_enabled = True

    def remove_subscriptions_for_research(self, research_id: str) -> None:
        """Remove all socket subscriptions for a completed research."""
        with self.__lock:
            removed = self.__socket_subscriptions.pop(research_id, None)
        if removed is not None:
            self.__log_info(
                f"Removed {len(removed)} subscription(s) for research {research_id}"
            )

    def __handle_connect(self, request):
        """Handle client connection"""
        username = session.get("username")
        if not username:
            self.__log_info(
                f"Rejected unauthenticated WebSocket connection from {request.sid}"
            )
            return False
        if not db_manager.is_user_connected(username):
            # Cookie is valid but the per-user DB engine isn't open yet (race vs first
            # XHR after page load, gunicorn worker restart, or idle eviction). Lazily
            # open it using the password the user authenticated with at login.
            session_id = session.get("session_id")
            password = (
                session_password_store.get_session_password(
                    username, session_id
                )
                if session_id
                else None
            )
            if not password:
                self.__log_info(
                    f"Rejected WebSocket connection for {username}: no active DB session and no stored password"
                )
                return False
            try:
                db_manager.open_user_database(username, password)
            except Exception as e:
                # Use __log_error (not __log_exception) so loguru cannot include
                # the `password` local in a diagnose=True traceback.
                self.__log_error(
                    f"Lazy DB open failed for {username} at WebSocket connect: {type(e).__name__}"
                )
                return False
        # Join a per-user room so user-scoped events (e.g. settings_changed,
        # which carries raw setting values including plaintext API keys) reach
        # only this user's own browser tabs and are never broadcast to every
        # connected client. Flask-SocketIO auto-removes the socket from the
        # room on disconnect.
        join_room(self.user_room(username))
        self.__log_info(f"Client connected: {request.sid} (user: {username})")
        return True

    def __handle_disconnect(self, request, reason: str):
        """Handle client disconnection"""
        try:
            self.__log_info(
                f"Client {request.sid} disconnected because: {reason}"
            )
            # Clean up subscriptions for this client.
            # __socket_subscriptions is keyed by research_id → set of sids,
            # so we iterate all entries and discard the disconnecting sid.
            with self.__lock:
                empty_keys = []
                for research_id, sids in self.__socket_subscriptions.items():
                    sids.discard(request.sid)
                    if not sids:
                        empty_keys.append(research_id)
                for key in empty_keys:
                    del self.__socket_subscriptions[key]
            self.__log_info(f"Removed subscription for client {request.sid}")

            # Clean up any thread-local database sessions that may have been
            # created during socket handler execution. This prevents file
            # descriptor leaks from unclosed SQLAlchemy sessions.
            try:
                from ...database.thread_local_session import (
                    cleanup_current_thread,
                )

                cleanup_current_thread()
            except ImportError:
                pass  # Module not available, skip cleanup
            except Exception:
                self.__log_exception(
                    "Error cleaning up thread session on disconnect"
                )
        except Exception as e:
            self.__log_exception(f"Error handling disconnect: {e}")

    def __handle_subscribe(self, data, request):
        """Handle client subscription to research updates."""
        research_id = data.get("research_id")
        if not research_id:
            return

        # Verify the connected user actually owns this research before
        # subscribing. The in-memory `_active_research` snapshot is keyed
        # only by research_id (no user tuple), so without this guard any
        # logged-in user could subscribe to any guessed/leaked research
        # UUID and receive its progress events. The per-user encrypted DB
        # is the ownership boundary: if the research row doesn't exist in
        # the user's DB, they don't own it.
        username = session.get("username")
        if not username or not self._user_owns_research(username, research_id):
            self.__log_info(
                f"Rejected subscribe from {request.sid}: user does not own research {research_id}"
            )
            return

        with self.__lock:
            if research_id not in self.__socket_subscriptions:
                self.__socket_subscriptions[research_id] = set()
            self.__socket_subscriptions[research_id].add(request.sid)
        self.__log_info(
            f"Client {request.sid} subscribed to research {research_id}"
        )

        # Send current status immediately if available in active research
        snapshot = get_active_research_snapshot(research_id)
        if snapshot is not None:
            progress = snapshot["progress"]
            latest_log = snapshot["log"][-1] if snapshot["log"] else None

            if latest_log:
                self.emit_socket_event(
                    f"progress_{research_id}",
                    {
                        "progress": progress,
                        "message": latest_log.get("message", "Processing..."),
                        "status": ResearchStatus.IN_PROGRESS,
                        "log_entry": latest_log,
                    },
                    room=request.sid,
                )

    @staticmethod
    def _user_owns_research(username: str, research_id: str) -> bool:
        """Return True if the given user owns this research / benchmark id.

        Used as the authorization boundary for WebSocket subscriptions —
        ownership is checked against the user's encrypted SQLite database,
        which is the per-user data partition. A static helper so unit
        tests can exercise the authz logic without standing up the
        singleton/Flask app.

        Recognizes both normal research (``ResearchHistory``, UUID id) and
        benchmark runs (``BenchmarkRun``, integer id) — the benchmark page
        subscribes with its ``BenchmarkRun.id``, which lives in the same
        per-user DB. Both checks stay scoped to the caller's own database,
        so no cross-user access is introduced.
        """
        try:
            from ...database.session_context import get_user_db_session
            from ...database.models import ResearchHistory

            with get_user_db_session(username) as db:
                if (
                    db.query(ResearchHistory.id)
                    .filter(ResearchHistory.id == research_id)
                    .first()
                    is not None
                ):
                    return True

                # Benchmark pages subscribe with their BenchmarkRun.id.
                # Recognize the user's own benchmark runs so the ownership
                # gate doesn't drop benchmark live progress (regression vs.
                # the removed cross-user broadcast). research_id stays a
                # string (never coerced to int — IDs are strings/UUIDs
                # repo-wide); SQLite applies numeric affinity to match the
                # Integer column. Only attempt this for numeric ids.
                if str(research_id).isdigit():
                    from ...database.models.benchmark import BenchmarkRun

                    return (
                        db.query(BenchmarkRun.id)
                        .filter(BenchmarkRun.id == research_id)
                        .first()
                        is not None
                    )
                return False
        except Exception:
            # Conservative: deny on any DB-open or query failure so a
            # transient infra error never silently widens authz.
            logger.opt(exception=True).warning(
                "Failed to verify research ownership for socket subscribe"
            )
            return False

    def __handle_unsubscribe(self, data, request):
        """Handle client unsubscribe from research updates."""
        research_id = (
            data.get("research_id") if isinstance(data, dict) else None
        )
        if not research_id:
            return

        # Symmetric with __handle_subscribe: require the caller to own the
        # research before mutating the per-research subscription set. The
        # practical impact of an unguarded unsubscribe is small (no data
        # exfiltration; subscribe is already guarded), but it keeps the
        # authz boundary consistent and avoids log spam from spoofed sids.
        username = session.get("username")
        if not username or not self._user_owns_research(username, research_id):
            self.__log_info(
                f"Rejected unsubscribe from {request.sid}: user does not own research {research_id}"
            )
            return

        with self.__lock:
            subs = self.__socket_subscriptions.get(research_id)
            if subs:
                subs.discard(request.sid)
                # Prune empty sets so the dict doesn't grow unbounded with
                # stale research_ids over long server runtimes.
                if not subs:
                    self.__socket_subscriptions.pop(research_id, None)
        self.__log_info(
            f"Client {request.sid} unsubscribed from research {research_id}"
        )

    def __handle_socket_error(self, e):
        """Handle Socket.IO errors"""
        self.__log_exception(f"Socket.IO error: {str(e)}")
        # Don't propagate exceptions to avoid crashing the server
        return False

    def __handle_default_error(self, e):
        """Handle unhandled Socket.IO errors"""
        self.__log_exception(f"Unhandled Socket.IO error: {str(e)}")
        # Don't propagate exceptions to avoid crashing the server
        return False

    def run(self, host: str, port: int, debug: bool = False) -> None:
        """
        Runs the SocketIO server.

        Args:
            host: The hostname to bind the server to.
            port: The port number to listen on.
            debug: Whether to run in debug mode. Defaults to False.

        """
        # Suppress Server header to prevent version information disclosure
        # This must be done before starting the server because Werkzeug adds
        # the header at the HTTP layer, not WSGI layer
        try:
            from werkzeug.serving import WSGIRequestHandler

            WSGIRequestHandler.version_string = lambda self: ""  # type: ignore[method-assign]
            logger.debug("Suppressed Server header for security")
        except ImportError:
            logger.warning(
                "Could not suppress Server header - werkzeug not found"
            )

        logger.info(f"Starting web server on {host}:{port} (debug: {debug})")
        self.__socketio.run(
            self.__app,  # Use the stored Flask app reference
            debug=debug,
            host=host,
            port=port,
            allow_unsafe_werkzeug=True,
            use_reloader=False,
        )
