"""
Encrypted database management using SQLCipher.
Handles per-user encrypted databases with browser-friendly authentication.
"""

import os
import threading
from pathlib import Path
from typing import Any, Dict, Optional, Tuple

from loguru import logger
from sqlalchemy import create_engine, event, text
from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session, sessionmaker
from sqlalchemy.pool import QueuePool, NullPool, StaticPool

from ..config.paths import get_data_directory, get_user_database_filename
from ..settings.env_registry import get_env_setting
from .sqlcipher_compat import get_sqlcipher_module
from .sqlcipher_utils import (
    set_sqlcipher_key,
    set_sqlcipher_rekey,
    apply_cipher_defaults_before_key,
    apply_sqlcipher_pragmas,
    apply_performance_pragmas,
    verify_sqlcipher_connection,
    create_database_salt,
    has_per_database_salt,
    get_key_from_password,
    get_sqlcipher_version,
    create_sqlcipher_connection,
)


class DatabaseManager:
    """Manages encrypted SQLCipher databases for each user."""

    def __init__(self):
        self.connections: Dict[str, Engine] = {}
        self._connections_lock = threading.RLock()
        # Track thread-specific engines for cleanup (key: (username, thread_id) tuple)
        self._thread_engines: Dict[Tuple[str, int], Engine] = {}
        self._thread_engine_lock = threading.Lock()
        self.data_dir = get_data_directory() / "encrypted_databases"
        self.data_dir.mkdir(parents=True, exist_ok=True)

        # Check SQLCipher availability
        self.has_encryption = self._check_encryption_available()

        # Determine pool class based on environment
        # Use StaticPool for testing to avoid locking issues
        self._use_static_pool = bool(os.environ.get("TESTING"))
        self._pool_class = StaticPool if self._use_static_pool else QueuePool

    def _get_pool_kwargs(self) -> Dict[str, Any]:
        """Get pool configuration kwargs based on pool type.

        StaticPool doesn't support pool_size or max_overflow,
        so we only include them for QueuePool.
        """
        if self._use_static_pool:
            return {}
        return {
            "pool_size": 10,
            "max_overflow": 30,  # Increase overflow so we don't run out of connections
            "pool_pre_ping": True,  # Validate connections before use (prevents stale connections)
        }

    def _is_valid_encryption_key(self, password: str) -> bool:
        """
        Check if the provided password is valid (not None, empty, or whitespace-only).

        Args:
            password: The password to check

        Returns:
            True if the password is valid, False otherwise
        """
        return password is not None and password.strip() != ""

    def is_user_connected(self, username: str) -> bool:
        """Check if a user has an active database connection.

        Thread-safe accessor for external callers.

        Args:
            username: The username to check

        Returns:
            True if the user has an active connection
        """
        with self._connections_lock:
            return username in self.connections

    def _check_encryption_available(self) -> bool:
        """Check if SQLCipher is available for encryption."""
        try:
            import os as os_module
            import tempfile

            # Test if SQLCipher actually works, not just if it imports
            with tempfile.NamedTemporaryFile(delete=False) as tmp:
                tmp_path = tmp.name

            try:
                # Try to create a test encrypted database
                sqlcipher_module = get_sqlcipher_module()
                sqlcipher = sqlcipher_module.dbapi2

                conn = sqlcipher.connect(tmp_path)
                try:
                    cursor = conn.cursor()
                    # Use creation_mode=True since we're creating a new test database
                    apply_cipher_defaults_before_key(cursor)
                    # Use centralized key setting
                    set_sqlcipher_key(cursor, "testpass")
                    # Apply post-key pragmas (kdf_iter for new DB)
                    apply_sqlcipher_pragmas(cursor, creation_mode=True)
                    apply_performance_pragmas(cursor)

                    # Check SQLCipher version
                    version = get_sqlcipher_version(cursor)
                    if version:
                        major = (
                            version.split(".")[0]
                            if "." in version
                            else version[0]
                        )
                        if major.isdigit() and int(major) < 4:
                            logger.warning(
                                f"SQLCipher version {version} detected. "
                                "Version 4.x+ is recommended for proper PRAGMA ordering."
                            )

                    cursor.close()
                    # Now use the connection for table operations
                    conn.execute("CREATE TABLE test (id INTEGER PRIMARY KEY)")
                    conn.execute("INSERT INTO test VALUES (1)")
                    result = conn.execute("SELECT * FROM test").fetchone()

                    if result != (1,):
                        raise Exception("SQLCipher encryption test failed")
                    logger.info(
                        "SQLCipher available and working - databases will be encrypted"
                    )
                    return True
                finally:
                    conn.close()
            except Exception as e:
                logger.warning(f"SQLCipher module found but not working: {e}")
                raise ImportError("SQLCipher not functional")
            finally:
                # Clean up test file
                try:
                    os_module.unlink(tmp_path)
                except OSError as e:
                    logger.debug(
                        f"Failed to clean up temp file {tmp_path}: {e}"
                    )

        except ImportError:
            # Check if user has explicitly allowed unencrypted databases.
            # Registry handles deprecated LDR_ALLOW_UNENCRYPTED fallback automatically.
            allow_unencrypted = get_env_setting(
                "bootstrap.allow_unencrypted", False
            )

            if not allow_unencrypted:
                logger.exception(
                    "SECURITY ERROR: SQLCipher is not installed!\n"
                    "Your databases will NOT be encrypted.\n"
                    "To fix this:\n"
                    "1. Install SQLCipher: sudo apt install sqlcipher libsqlcipher-dev\n"
                    "2. Reinstall project: pdm install\n"
                    "Or use Docker with SQLCipher pre-installed.\n\n"
                    "To explicitly allow unencrypted databases (NOT RECOMMENDED):\n"
                    "export LDR_BOOTSTRAP_ALLOW_UNENCRYPTED=true"
                )
                raise RuntimeError(
                    "SQLCipher not available. Set LDR_BOOTSTRAP_ALLOW_UNENCRYPTED=true to proceed without encryption (NOT RECOMMENDED)"
                )
            else:
                logger.warning(
                    "WARNING: Running with UNENCRYPTED databases!\n"
                    "This means:\n"
                    "- Passwords don't protect data access\n"
                    "- API keys are stored in plain text\n"
                    "- Anyone with file access can read all data\n"
                    "Install SQLCipher for secure operation!"
                )
                return False

    def _get_user_db_path(self, username: str) -> Path:
        """Get the path for a user's encrypted database."""
        return self.data_dir / get_user_database_filename(username)

    def _apply_pragmas(self, connection, connection_record):
        """Apply pragmas for optimal performance."""
        # Check if this is SQLCipher or regular SQLite
        is_encrypted = self.has_encryption

        # Use centralized performance pragma application

        apply_performance_pragmas(connection)

        # SQLCipher-specific pragmas
        if is_encrypted:
            from .sqlcipher_utils import get_sqlcipher_settings

            settings = get_sqlcipher_settings()
            pragmas = [
                f"PRAGMA kdf_iter = {settings['kdf_iterations']}",
                f"PRAGMA cipher_page_size = {settings['page_size']}",
            ]
            for pragma in pragmas:
                try:
                    connection.execute(pragma)
                except Exception as e:
                    logger.debug(f"Could not apply pragma '{pragma}': {e}")
        else:
            # Regular SQLite pragma
            try:
                connection.execute(
                    "PRAGMA mmap_size = 268435456"
                )  # 256MB memory mapping
            except Exception as e:
                logger.debug(f"Could not apply mmap_size pragma: {e}")

    @staticmethod
    def _make_sqlcipher_connection(
        db_path: Path,
        password: str,
        isolation_level: Optional[str] = "IMMEDIATE",
        check_same_thread: bool = False,
    ) -> Any:
        """Create a properly initialized SQLCipher connection.

        Follows the canonical SQLCipher initialization order: set key,
        apply cipher pragmas, verify, then apply performance pragmas.
        Cipher pragmas (page size, HMAC algorithm, KDF iterations) must
        be configured before the first query (verification) because that
        query triggers page decryption with the active cipher settings.

        Args:
            db_path: Path to the database file
            password: The database encryption passphrase
            isolation_level: SQLite isolation level (``""`` for deferred
                transactions, ``None`` for autocommit)
            check_same_thread: SQLite check_same_thread flag

        Returns:
            A raw ``sqlcipher3`` connection ready for use.

        Raises:
            ValueError: If the database key cannot be verified.
        """
        sqlcipher3 = get_sqlcipher_module()
        conn = sqlcipher3.connect(
            str(db_path),
            isolation_level=isolation_level,
            check_same_thread=check_same_thread,
        )
        cursor = conn.cursor()

        try:
            set_sqlcipher_key(cursor, password, db_path=db_path)
            apply_sqlcipher_pragmas(cursor, creation_mode=False)

            if not verify_sqlcipher_connection(cursor):
                raise ValueError("Failed to verify database key")

            apply_performance_pragmas(cursor)
        except Exception:
            try:
                cursor.close()
            except Exception:  # noqa: BLE001
                logger.warning(
                    "Failed to close cursor during cleanup", exc_info=True
                )
            conn.close()
            raise

        cursor.close()
        return conn

    def create_user_database(self, username: str, password: str) -> Engine:
        """Create a new encrypted database for a user."""

        # Validate the encryption key
        if not self._is_valid_encryption_key(password):
            logger.error(
                f"Invalid encryption key for user {username}: password is None or empty"
            )
            raise ValueError(
                "Invalid encryption key: password cannot be None or empty"
            )

        db_path = self._get_user_db_path(username)

        if db_path.exists():
            raise ValueError(f"Database already exists for user {username}")

        # Create connection string - use regular SQLite when SQLCipher not available
        if self.has_encryption:
            # Create directory if it doesn't exist
            db_path.parent.mkdir(parents=True, exist_ok=True)

            # Create per-database salt for new databases (v2 security improvement)
            create_database_salt(db_path)
            logger.info(f"Created per-database salt for {username}")

            # Pre-derive key before closures to avoid capturing plaintext password
            hex_key = get_key_from_password(password, db_path=db_path).hex()

            # Create database structure using raw SQLCipher outside SQLAlchemy
            try:
                conn = create_sqlcipher_connection(
                    db_path,
                    password=password,
                    creation_mode=True,
                    connect_kwargs={
                        "isolation_level": "IMMEDIATE",
                        "check_same_thread": False,
                    },
                )
                try:
                    # Get the CREATE TABLE statements from SQLAlchemy models
                    from sqlalchemy.dialects import sqlite
                    from sqlalchemy.schema import CreateTable

                    from .models import Base

                    # Create tables one by one
                    sqlite_dialect = sqlite.dialect()
                    for table in Base.metadata.sorted_tables:
                        if table.name != "users":
                            create_sql = str(
                                CreateTable(table).compile(
                                    dialect=sqlite_dialect
                                )
                            )
                            logger.debug(f"Creating table {table.name}")
                            conn.execute(create_sql)

                    conn.commit()
                finally:
                    conn.close()

                logger.info(
                    f"Database structure created successfully for {username}"
                )

            except Exception:
                logger.exception("Error creating database structure")
                # Cleanup partial DB file on failure
                if db_path.exists():
                    db_path.unlink(missing_ok=True)
                raise

            # Small delay to ensure file is fully written
            import time

            time.sleep(0.1)

            # Now create SQLAlchemy engine using custom connection creator
            def create_engine_connection():
                """Create a properly initialized SQLCipher connection."""
                return create_sqlcipher_connection(
                    db_path,
                    hex_key=hex_key,
                    creation_mode=False,
                    connect_kwargs={
                        "isolation_level": "IMMEDIATE",
                        "check_same_thread": False,
                    },
                )

            # Create engine with custom creator function and optimized cache
            engine = create_engine(
                "sqlite://",
                creator=create_engine_connection,
                poolclass=self._pool_class,
                echo=False,
                query_cache_size=1000,
                **self._get_pool_kwargs(),
            )
        else:
            logger.warning(
                f"SQLCipher not available - creating UNENCRYPTED database for user {username}"
            )
            # Fall back to regular SQLite with query cache
            engine = create_engine(
                f"sqlite:///{db_path}",
                connect_args={"check_same_thread": False, "timeout": 30},
                poolclass=self._pool_class,
                echo=False,
                query_cache_size=1000,
                **self._get_pool_kwargs(),
            )

            # For unencrypted databases, just apply pragmas
            event.listen(engine, "connect", self._apply_pragmas)

        # Tables have already been created using raw SQLCipher above
        # No need to create them again with SQLAlchemy

        # Store connection
        with self._connections_lock:
            self.connections[username] = engine

        # Initialize database tables using centralized initialization
        from .initialize import initialize_database

        try:
            # Create a session for settings initialization
            Session = sessionmaker(bind=engine)
            with Session() as session:
                initialize_database(engine, session)
        except Exception as e:
            logger.warning(f"Could not initialize database fully: {e}")
            # Still continue - basic tables were created above

        logger.info(f"Created encrypted database for user {username}")
        return engine

    def open_user_database(
        self, username: str, password: str
    ) -> Optional[Engine]:
        """Open an existing encrypted database for a user."""

        # Validate the encryption key
        if not self._is_valid_encryption_key(password):
            logger.error(
                f"Invalid encryption key when opening database for user {username}: password is None or empty"
            )
            # TODO: Fix the root cause - research threads are not getting the correct password
            logger.error(
                "TODO: This usually means the research thread is not receiving the user's "
                "password for database encryption. Need to ensure password is passed from "
                "the main thread to research threads."
            )
            raise ValueError(
                "Invalid encryption key: password cannot be None or empty"
            )

        # Check if already open
        with self._connections_lock:
            if username in self.connections:
                return self.connections[username]

        db_path = self._get_user_db_path(username)

        # Prevent timing attacks: always derive key before checking file existence
        # This ensures both existing and non-existent users take the same amount of time,
        # preventing username enumeration via timing analysis.
        # Pre-derive key before closures to avoid capturing plaintext password
        hex_key = get_key_from_password(password, db_path=db_path).hex()

        if not db_path.exists():
            logger.error(f"No database found for user {username}")
            return None

        # Warn if this is a legacy database without per-database salt
        if self.has_encryption and not has_per_database_salt(db_path):
            logger.warning(
                f"Database for user '{username}' uses the legacy shared salt "
                f"(deprecated). For improved security, consider creating a new "
                f"account to get a per-database salt. Legacy databases remain "
                f"fully functional but are less resistant to multi-target attacks."
            )

        # Create connection string - use regular SQLite when SQLCipher not available
        if self.has_encryption:

            def create_open_connection():
                """Create a properly initialized SQLCipher connection."""
                return create_sqlcipher_connection(
                    db_path,
                    hex_key=hex_key,
                    creation_mode=False,
                    connect_kwargs={
                        "isolation_level": "IMMEDIATE",
                        "check_same_thread": False,
                    },
                )

            # Create engine with custom creator function and optimized cache
            engine = create_engine(
                "sqlite://",
                creator=create_open_connection,
                poolclass=self._pool_class,
                echo=False,
                query_cache_size=1000,
                **self._get_pool_kwargs(),
            )
        else:
            logger.warning(
                f"SQLCipher not available - opening UNENCRYPTED database for user {username}"
            )
            # Fall back to regular SQLite (no password protection!)
            engine = create_engine(
                f"sqlite:///{db_path}",
                connect_args={"check_same_thread": False, "timeout": 30},
                poolclass=self._pool_class,
                echo=False,
                query_cache_size=1000,
                **self._get_pool_kwargs(),
            )

            # For unencrypted databases, just apply pragmas
            event.listen(engine, "connect", self._apply_pragmas)

        try:
            # Test connection by running a simple query
            with engine.connect() as conn:
                conn.execute(text("SELECT 1"))

            # Store connection
            with self._connections_lock:
                self.connections[username] = engine

            # Run database initialization (creates missing tables and runs migrations)
            from .initialize import initialize_database

            try:
                initialize_database(engine)
            except Exception as e:
                logger.warning(f"Could not run migrations for {username}: {e}")

            logger.info(f"Opened encrypted database for user {username}")
            return engine

        except Exception:
            logger.exception(f"Failed to open database for user {username}")
            engine.dispose()
            return None

    def get_session(self, username: str) -> Optional[Session]:
        """Create a new session for a user's database."""
        with self._connections_lock:
            if username not in self.connections:
                # Use debug level for this common scenario to reduce log noise
                logger.debug(f"No open database for user {username}")
                return None
            engine = self.connections[username]
            # Create session inside lock to prevent race with close_user_database()
            SessionLocal = sessionmaker(bind=engine)
            return SessionLocal()

    def close_user_database(self, username: str):
        """Close a user's database connection."""
        with self._connections_lock:
            if username in self.connections:
                self.connections[username].dispose()
                del self.connections[username]
                logger.info(f"Closed database for user {username}")

        # Also cleanup any thread engines for this user
        self.cleanup_thread_engines(username=username)

    def check_database_integrity(self, username: str) -> bool:
        """Check integrity of a user's encrypted database."""
        with self._connections_lock:
            if username not in self.connections:
                return False
            engine = self.connections[username]

        try:
            with engine.connect() as conn:
                # Quick integrity check
                result = conn.execute(text("PRAGMA quick_check"))
                if result.fetchone()[0] != "ok":
                    return False

                # SQLCipher integrity check
                result = conn.execute(text("PRAGMA cipher_integrity_check"))
                # If this returns any rows, there are HMAC failures
                failures = list(result)
                if failures:
                    logger.error(
                        f"Integrity check failed for {username}: {len(failures)} HMAC failures"
                    )
                    return False

                return True

        except Exception:
            logger.exception(f"Integrity check error for user: {username}")
            return False

    def change_password(
        self, username: str, old_password: str, new_password: str
    ) -> bool:
        """Change the encryption password for a user's database.

        This rekeys the SQLCipher database — no separate auth-DB
        password-hash update is needed because passwords are never
        stored.  Login verification is done by attempting decryption.
        """
        if not self.has_encryption:
            logger.warning(
                "Cannot change password - SQLCipher not available (databases are unencrypted)"
            )
            return False

        db_path = self._get_user_db_path(username)

        if not db_path.exists():
            return False

        try:
            # Close existing connection if any
            self.close_user_database(username)

            # Open with old password
            engine = self.open_user_database(username, old_password)
            if not engine:
                return False

            # Rekey the database (only works with SQLCipher)
            with engine.connect() as conn:
                # Use centralized rekey function
                set_sqlcipher_rekey(conn, new_password, db_path=db_path)

            logger.info(f"Password changed for user {username}")
            return True

        except Exception:
            logger.exception(f"Failed to change password for user: {username}")
            return False
        finally:
            # Close the connection
            self.close_user_database(username)

    def user_exists(self, username: str) -> bool:
        """Check if a user exists in the auth database."""
        from .auth_db import auth_db_session
        from .models.auth import User

        with auth_db_session() as session:
            user = session.query(User).filter_by(username=username).first()
            return user is not None

    def get_memory_usage(self) -> Dict[str, Any]:
        """Get memory usage statistics."""
        with self._connections_lock:
            num_connections = len(self.connections)
        return {
            "active_connections": num_connections,
            "thread_engines": len(self._thread_engines),
            "active_sessions": 0,  # Sessions are created on-demand, not tracked
            "estimated_memory_mb": (num_connections + len(self._thread_engines))
            * 3.5,  # ~3.5MB per connection
        }

    def cleanup_thread_engines(
        self, username: str = None, thread_id: int = None
    ):
        """
        Clean up thread-specific engines.

        Args:
            username: If provided, cleanup engines for this user only
            thread_id: If provided, cleanup engines for this thread only
                      If neither provided, cleanup current thread's engines
        """
        if thread_id is None and username is None:
            thread_id = threading.get_ident()

        with self._thread_engine_lock:
            keys_to_remove = []
            for key in self._thread_engines:
                # Key is a tuple: (username, thread_id)
                key_user, key_thread = key

                should_remove = False
                if username and thread_id:
                    should_remove = (
                        key_user == username and key_thread == thread_id
                    )
                elif username:
                    should_remove = key_user == username
                elif thread_id:
                    should_remove = key_thread == thread_id

                if should_remove:
                    keys_to_remove.append(key)

            for key in keys_to_remove:
                try:
                    self._thread_engines[key].dispose()
                    logger.debug(f"Disposed thread engine: {key}")
                except Exception as e:
                    logger.warning(f"Error disposing thread engine {key}: {e}")
                del self._thread_engines[key]

    def cleanup_all_thread_engines(self):
        """Clean up all thread-specific engines (for shutdown)."""
        with self._thread_engine_lock:
            count = len(self._thread_engines)
            for key, engine in list(self._thread_engines.items()):
                try:
                    engine.dispose()
                except Exception as e:
                    logger.warning(f"Error disposing thread engine {key}: {e}")
            self._thread_engines.clear()
            logger.info(f"All thread engines disposed ({count} total)")

    def create_thread_safe_session_for_metrics(
        self, username: str, password: str
    ):
        """
        Create a new database session safe for use in background threads.
        This is specifically for metrics/logging - NOT for settings or user data.

        Args:
            username: The username
            password: The user's password (encryption key)

        Returns:
            A SQLAlchemy session that can be used in the current thread

        IMPORTANT: This should ONLY be used for:
        - Writing token metrics
        - Writing search metrics
        - Writing logs

        DO NOT use this for:
        - Reading/writing settings
        - Modifying user data
        - Any operation that should be synchronized with user requests
        """
        db_path = self._get_user_db_path(username)

        # Warn if this is a legacy database without per-database salt
        if self.has_encryption and not has_per_database_salt(db_path):
            logger.warning(
                f"Database for user '{username}' uses a legacy shared salt "
                f"(deprecated). For improved security, consider creating a new "
                f"account to get a per-database salt. Legacy databases remain "
                f"fully functional but are less resistant to multi-target attacks."
            )

        # Prevent timing attacks: always derive key before checking file existence
        # This ensures both existing and non-existent users take the same amount of time,
        # preventing username enumeration via timing analysis.
        # Pre-derive key before closures
        hex_key = get_key_from_password(password, db_path=db_path).hex()

        if not db_path.exists():
            raise ValueError(f"No database found for user {username}")

        # Create thread-specific key for engine reuse (tuple is cleaner than string)
        thread_id = threading.get_ident()
        engine_key = (username, thread_id)

        # Check for existing engine for this thread/user combo
        with self._thread_engine_lock:
            if engine_key in self._thread_engines:
                engine = self._thread_engines[engine_key]
                # Verify engine is still valid
                try:
                    with engine.connect() as conn:
                        conn.execute(text("SELECT 1"))
                    # Engine is valid, create session from it
                    Session = sessionmaker(bind=engine)
                    return Session()
                except Exception as e:
                    # Engine is stale, remove it and create new one
                    logger.debug(
                        f"Stale thread engine for {engine_key}, recreating: {e}"
                    )
                    try:
                        engine.dispose()
                    except Exception as dispose_error:
                        logger.debug(
                            f"Error disposing stale engine: {dispose_error}"
                        )
                    del self._thread_engines[engine_key]

        # Create a thread-local engine
        if self.has_encryption:

            def create_thread_connection():
                """Create a SQLCipher connection for this thread."""
                try:
                    return create_sqlcipher_connection(
                        db_path,
                        hex_key=hex_key,
                        creation_mode=False,
                        connect_kwargs={"check_same_thread": False},
                    )
                except Exception:
                    logger.exception(
                        f"Failed to create thread connection for {username}"
                    )
                    raise

            engine = create_engine(
                "sqlite://",
                creator=create_thread_connection,
                poolclass=NullPool,  # Important: no connection pooling for threads
                echo=False,
            )
        else:
            # Unencrypted fallback
            logger.warning("Creating unencrypted thread session for metrics")
            engine = create_engine(
                f"sqlite:///{db_path}",
                poolclass=NullPool,
                echo=False,
            )

        # Store engine for reuse within this thread
        with self._thread_engine_lock:
            self._thread_engines[engine_key] = engine

        # Create session
        Session = sessionmaker(bind=engine)
        return Session()


# Global instance
db_manager = DatabaseManager()
