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

import threading
from typing import Optional, Dict, Tuple
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:
            # Verify it's still valid
            try:
                # Simple query to test connection
                self._local.session.execute("SELECT 1")
                return self._local.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."""
        thread_id = threading.get_ident()

        if hasattr(self._local, "session") and self._local.session:
            try:
                self._local.session.close()
                logger.debug(f"Thread {thread_id}: Closed database session")
            except Exception as e:
                logger.exception(
                    f"Thread {thread_id}: Error closing session: {e}"
                )
            finally:
                self._local.session = 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_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."""
    thread_session_manager.cleanup_thread()


# 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
