"""
Thread-local database session management.
Each thread gets its own database session that persists for the thread's lifetime.
"""

import functools
import threading
from contextlib import ContextDecorator
from typing import Optional, Dict, Tuple
from sqlalchemy import text
from sqlalchemy.exc import PendingRollbackError
from sqlalchemy.orm import Session
from loguru import logger

from .encrypted_db import db_manager


class ThreadLocalSessionManager:
    """
    Manages database sessions per thread.
    Each thread gets its own session that is reused throughout the thread's lifetime.
    """

    def __init__(self):
        # Thread-local storage for sessions
        self._local = threading.local()
        # Track credentials per thread ID (for cleanup)
        self._thread_credentials: Dict[int, Tuple[str, str]] = {}
        self._lock = threading.Lock()

    def get_session(self, username: str, password: str) -> Optional[Session]:
        """
        Get or create a database session for the current thread.

        The session is created once per thread and reused for all subsequent calls.
        This avoids the expensive SQLCipher decryption on every database access.
        """
        thread_id = threading.get_ident()

        # Check if we already have a session for this thread
        if hasattr(self._local, "session") and self._local.session:
            # SECURITY: ensure cached session belongs to the requesting user
            if getattr(self._local, "username", None) != username:
                logger.warning(
                    f"Thread {thread_id}: Session username mismatch "
                    f"(cached={self._local.username!r}, requested={username!r}), "
                    "clearing stale cross-user session"
                )
                self._cleanup_thread_session()
            else:
                # Verify it's still valid
                try:
                    self._local.session.execute(text("SELECT 1"))
                    # Under DEFERRED isolation the validation SELECT opens
                    # a transaction that holds a SHARED lock on SQLite
                    # until an explicit commit/rollback. A long-lived
                    # thread-local session reused across requests would
                    # keep that lock held and block the first writer.
                    # Roll it back so subsequent callers start fresh.
                    self._local.session.rollback()
                    return self._local.session
                except PendingRollbackError:
                    # Session has a pending rollback (e.g. from a previous database lock error).
                    # Attempt rollback to recover without destroying the session.
                    logger.debug(
                        f"Thread {thread_id}: PendingRollbackError, attempting rollback recovery"
                    )
                    try:
                        self._local.session.rollback()
                        self._local.session.execute(text("SELECT 1"))
                        self._local.session.rollback()
                        return self._local.session
                    except Exception:
                        logger.warning(
                            f"Thread {thread_id}: Rollback recovery failed, creating new session"
                        )
                        self._cleanup_thread_session()
                except Exception:
                    # Session is invalid, will create a new one
                    logger.debug(
                        f"Thread {thread_id}: Existing session invalid, creating new one"
                    )
                    self._cleanup_thread_session()

        # Create new session for this thread
        logger.debug(
            f"Thread {thread_id}: Creating new database session for user {username}"
        )

        # Ensure database is open. open_user_database returns None for
        # credential failures and raises DatabaseInitializationError when
        # the schema can't be initialised; from a worker-thread caller
        # both mean "no usable session right now" so collapse them.
        from .encrypted_db import DatabaseInitializationError

        try:
            engine = db_manager.open_user_database(username, password)
        except DatabaseInitializationError:
            # ``logger.warning`` (no traceback) rather than
            # ``logger.exception``: ``password`` is a live local in this
            # frame, so rendering a traceback under ``diagnose=True``
            # would dump the plaintext SQLCipher master password
            # (unrecoverable — TRUST.md §5). The redacted failure detail
            # is already logged at the raise site in
            # ``open_user_database`` (#4182).
            logger.warning(
                f"Thread {thread_id}: database init failed for user {username}"
            )
            return None
        if not engine:
            logger.error(
                f"Thread {thread_id}: Failed to open database for user {username}"
            )
            return None

        # Create session for this thread
        session = db_manager.create_thread_safe_session_for_metrics(
            username, password
        )
        if not session:
            logger.error(
                f"Thread {thread_id}: Failed to create session for user {username}"
            )
            return None

        # Store in thread-local storage
        self._local.session = session
        self._local.username = username

        # Track credentials for cleanup
        with self._lock:
            self._thread_credentials[thread_id] = (username, password)

        return session

    def get_current_session(self) -> Optional[Session]:
        """Get the current thread's session if it exists."""
        if hasattr(self._local, "session"):
            return self._local.session
        return None

    def _cleanup_thread_session(self):
        """Clean up the current thread's session.

        Sessions are bound to the shared per-user QueuePool engine, so
        closing the session returns its connection to the pool — there
        are no per-thread engines to dispose.
        """
        thread_id = threading.get_ident()

        if hasattr(self._local, "session") and self._local.session:
            try:
                self._local.session.rollback()
            except Exception:
                logger.warning(
                    f"Thread {thread_id}: Error rolling back session during cleanup"
                )
            try:
                self._local.session.close()
                logger.debug(f"Thread {thread_id}: Closed database session")
            except Exception:
                logger.warning(f"Thread {thread_id}: Error closing session")
            finally:
                self._local.session = None

        if hasattr(self._local, "username"):
            self._local.username = None

        # Remove from tracking
        with self._lock:
            self._thread_credentials.pop(thread_id, None)

    def cleanup_thread(self, thread_id: Optional[int] = None):
        """
        Clean up session for a specific thread or current thread.
        Called when a thread is finishing.
        """
        if thread_id is None:
            thread_id = threading.get_ident()

        # If it's the current thread, we can clean up directly
        if thread_id == threading.get_ident():
            self._cleanup_thread_session()
        else:
            # For other threads, just remove from tracking
            # The thread-local storage will be cleaned up when the thread ends
            with self._lock:
                self._thread_credentials.pop(thread_id, None)

    def cleanup_dead_threads(self):
        """Remove credential entries for threads that are no longer alive.

        Handles the abnormal case of threads that died without triggering
        their cleanup handler. Uses threading.enumerate() to identify
        alive threads and removes credential entries for dead ones.
        Sessions on dead threads are garbage-collected normally and
        their connections return to the shared per-user QueuePool.

        Called from:
        - processor_v2.py: every ~60s in the queue loop
        - app_factory.py: in teardown_appcontext
        - connection_cleanup.py: in cleanup_idle_connections (every ~300s)
        """
        alive_ids = {t.ident for t in threading.enumerate()}
        with self._lock:
            dead_ids = [
                tid for tid in self._thread_credentials if tid not in alive_ids
            ]
            for tid in dead_ids:
                del self._thread_credentials[tid]
        if dead_ids:
            logger.debug(f"Swept {len(dead_ids)} dead thread credential(s)")

    def cleanup_all(self):
        """Clean up all tracked sessions (for shutdown)."""
        with self._lock:
            thread_ids = list(self._thread_credentials.keys())

        for thread_id in thread_ids:
            self.cleanup_thread(thread_id)


# Global instance
thread_session_manager = ThreadLocalSessionManager()


def get_metrics_session(username: str, password: str) -> Optional[Session]:
    """
    Get a database session for metrics operations in the current thread.
    The session is created once and reused for the thread's lifetime.

    Note: This specifically uses create_thread_safe_session_for_metrics internally
    and should only be used for metrics-related database operations.
    """
    return thread_session_manager.get_session(username, password)


def get_current_thread_session() -> Optional[Session]:
    """Get the current thread's session if it exists."""
    return thread_session_manager.get_current_session()


def cleanup_current_thread():
    """Clean up the current thread's database session and cached credentials.

    Also clears any passwords cached by ``metrics_writer`` on this thread —
    pooled worker threads must not retain plaintext credentials across tasks.
    """
    thread_session_manager.cleanup_thread()
    try:
        from .thread_metrics import metrics_writer

        metrics_writer.clear_passwords()
    except Exception:
        logger.debug(
            "cleanup_current_thread: error clearing metrics_writer passwords",
            exc_info=True,
        )


def cleanup_dead_threads():
    """Sweep dead-thread credential entries from the session manager."""
    try:
        thread_session_manager.cleanup_dead_threads()
    except Exception:
        logger.warning("Dead-thread session sweep failed")


class _ThreadCleanup(ContextDecorator):
    """Context manager / decorator for thread-local resource cleanup."""

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        try:
            cleanup_current_thread()
        except Exception:
            logger.debug(
                "thread_cleanup: error during DB session cleanup",
                exc_info=True,
            )
        try:
            from ..config.thread_settings import clear_settings_context

            clear_settings_context()
        except Exception:
            logger.debug(
                "thread_cleanup: error clearing settings context",
                exc_info=True,
            )
        try:
            from ..utilities.thread_context import clear_search_context

            clear_search_context()
        except Exception:
            logger.debug(
                "thread_cleanup: error clearing search context",
                exc_info=True,
            )
        try:
            # Defense-in-depth audit hook (PEP 578). Leaving an active
            # EgressContext on a pooled worker thread would let the
            # next, unrelated task inherit the previous run's scope.
            from ..security.egress.audit_hook import clear_active_context

            clear_active_context()
        except Exception:
            logger.debug(
                "thread_cleanup: error clearing egress audit context",
                exc_info=True,
            )
        return False


def thread_cleanup(func=None):
    """Ensure all thread-local resources are cleaned up when a function or block exits.

    Works as a bare decorator, a decorator factory, or a context manager::

        @thread_cleanup
        def worker(): ...

        @thread_cleanup()
        def worker(): ...

        with thread_cleanup():
            ...

        executor.submit(thread_cleanup(func), arg)
    """
    if func is not None:

        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            with _ThreadCleanup():
                return func(*args, **kwargs)

        return wrapper
    return _ThreadCleanup()


# Context manager for automatic cleanup
class ThreadSessionContext:
    """
    Context manager that ensures thread session is cleaned up.
    Usage:
        with ThreadSessionContext(username, password) as session:
            # Use session
    """

    def __init__(self, username: str, password: str):
        self.username = username
        self.password = password
        self.session = None

    def __enter__(self) -> Optional[Session]:
        self.session = get_metrics_session(self.username, self.password)
        return self.session

    def __exit__(self, exc_type, exc_val, exc_tb):
        # Don't cleanup here - let the thread keep its session
        # Only cleanup when thread ends
        pass
