"""
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"))
                    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"))
                        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
        engine = db_manager.open_user_database(username, password)
        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 and engine."""
        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

        # Clean up the thread engine too
        if hasattr(self._local, "username") and self._local.username:
            db_manager.cleanup_thread_engines(
                username=self._local.username, thread_id=thread_id
            )
            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.

        Companion to encrypted_db.cleanup_dead_thread_engines().
        Removes credentials and session tracking for threads no longer alive.

        This is separate from cleanup_current_thread() which handles the
        normal case (thread cleans up its own session on exit via
        @thread_cleanup decorator).

        cleanup_dead_threads() handles the abnormal case: threads that died
        without triggering their cleanup handler. It uses
        threading.enumerate() to identify alive threads and removes entries
        for dead ones.

        Called from:
        - processor_v2.py: every ~60s in the queue loop
        - app_factory.py: in teardown_appcontext (engine sweep rate-limited)
        - 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)

        # Also cleanup any remaining thread engines
        db_manager.cleanup_all_thread_engines()


# 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."""
    thread_session_manager.cleanup_thread()


def cleanup_dead_threads():
    """Sweep dead-thread entries from both session manager and engine dict.

    Each call is independently guarded so one failure does not block the other.
    Engine sweep is rate-limited internally (at most once per 60 s).
    """
    try:
        db_manager.maybe_sweep_dead_engines()
    except Exception:
        logger.warning("Dead-thread engine sweep failed")
    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,
            )
        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
