"""
Database session context manager and decorator for encrypted databases.
Ensures all database access has proper encryption context.
"""

import functools
from contextlib import contextmanager
from typing import Callable, Optional

from flask import (
    g,
    has_app_context,
    has_request_context,
    jsonify,
    session as flask_session,
)
from loguru import logger
from sqlalchemy.orm import Session

from ..utilities.thread_context import get_search_context
from .encrypted_db import db_manager

# Placeholder password used when accessing unencrypted databases.
# This should only be used when LDR_ALLOW_UNENCRYPTED=true is set.
UNENCRYPTED_DB_PLACEHOLDER = "unencrypted-mode"


class DatabaseSessionError(Exception):
    """Raised when database session cannot be established."""

    pass


def safe_rollback(session: Session, context: str = "") -> None:
    """Roll back the session, swallowing and logging any rollback failure.

    SQLAlchemy requires explicit rollback after a failed flush/commit before
    the session is usable again. Skipping it leaves the session in
    PendingRollbackError state and every subsequent ORM operation cascades.

    This helper exists so call sites can recover the session in one line
    without repeating the try/except/log boilerplate at every except handler.
    ``context`` is included in the error log so failed rollbacks can be
    traced back to the call site.
    """
    try:
        session.rollback()
    except Exception:
        if context:
            logger.exception(f"Failed to rollback session: {context}")
        else:
            logger.exception("Failed to rollback session")


def get_g_db_session() -> Optional[Session]:
    """Lazily create and cache a DB session on Flask g for the current request.

    Follows Flask's recommended pattern for lazy resource creation:
    only check out a pool connection when a route actually needs one.
    Returns None if no user is authenticated or DB is not connected.
    """
    if (
        has_app_context()
        and hasattr(g, "db_session")
        and g.db_session is not None
    ):
        return g.db_session
    username = getattr(g, "current_user", None) if has_app_context() else None
    if not username or not db_manager.is_user_connected(username):
        return None
    try:
        session = db_manager.get_session(username)
        g.db_session = session
        return session
    except Exception:
        logger.exception(f"Error lazily creating session for {username}")
        return None


@contextmanager
def get_user_db_session(
    username: Optional[str] = None, password: Optional[str] = None
):
    """
    Context manager that ensures proper database session with encryption.
    Now uses thread-local sessions for better performance.

    Args:
        username: Username (if not provided, gets from Flask session)
        password: Password for encrypted database (required for first access)

    Yields:
        Database session for the user

    Raises:
        DatabaseSessionError: If session cannot be established
    """
    # Import here to avoid circular imports
    from .thread_local_session import get_metrics_session
    from .session_passwords import session_password_store

    session = None
    needs_close = False

    try:
        # Get username from Flask session if not provided (only in Flask context)
        if not username and has_request_context():
            username = flask_session.get("username")

        if not username:
            raise DatabaseSessionError("No authenticated user")

        # First, try to get/create a session via Flask g (best performance)
        if has_app_context():
            cached = get_g_db_session()
            if cached:
                session = cached
                needs_close = False

        if not session:
            # Get password if not provided
            if not password and has_app_context():
                # Try to get from g (works with app context)
                if hasattr(g, "user_password"):
                    password = g.user_password
                    logger.debug(
                        f"Got password from g.user_password for {username}"
                    )

            # Try session password store (requires request context for flask_session)
            if not password and has_request_context():
                session_id = flask_session.get("session_id")
                if session_id:
                    logger.debug(
                        f"Trying session password store for {username}"
                    )
                    password = session_password_store.get_session_password(
                        username, session_id
                    )
                    if password:
                        logger.debug(
                            f"Got password from session store for {username}"
                        )
                    else:
                        logger.debug(
                            f"No password in session store for {username}"
                        )

            # Try thread context (for background threads)
            if not password:
                thread_context = get_search_context()
                if thread_context and thread_context.get("user_password"):
                    password = thread_context["user_password"]
                    logger.debug(
                        f"Got password from thread context for {username}"
                    )

            if not password and db_manager.has_encryption:
                raise DatabaseSessionError(
                    f"Encrypted database for {username} requires password"
                )
            elif not password:
                logger.warning(
                    f"Accessing unencrypted database for {username} - "
                    "ensure this is intentional (LDR_ALLOW_UNENCRYPTED=true)"
                )
                password = UNENCRYPTED_DB_PLACEHOLDER

            # Use thread-local session (will reuse existing or create new)
            session = get_metrics_session(username, password)
            if not session:
                raise DatabaseSessionError(
                    f"Could not establish session for {username}"
                )
            # Thread-local sessions are managed by the thread, don't close them
            needs_close = False

            # Store the password we successfully used
            if password and has_app_context():
                g.user_password = password

        yield session

    finally:
        # Only close if we created a new session (which we don't anymore)
        if session and needs_close:
            try:
                session.close()
            except Exception:
                logger.debug("Failed to close session during cleanup")


def with_user_database(func: Callable) -> Callable:
    """
    Decorator that ensures function has access to user's database.
    Injects 'db_session' as first argument to the decorated function.

    Usage:
        @with_user_database
        def get_user_settings(db_session, setting_key):
            return db_session.query(Setting).filter_by(key=setting_key).first()
    """

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        # Check if username/password provided in kwargs
        username = kwargs.pop("_username", None)
        password = kwargs.pop("_password", None)

        with get_user_db_session(username, password) as db_session:
            return func(db_session, *args, **kwargs)

    return wrapper


def ensure_db_session(view_func: Callable) -> Callable:
    """
    Flask view decorator that ensures database session is available.
    Sets g.db_session for use in the request.

    Usage:
        @app.route('/my-route')
        @ensure_db_session
        def my_view():
            # g.db_session is available here
            settings = g.db_session.query(Setting).all()
    """

    @functools.wraps(view_func)
    def wrapper(*args, **kwargs):
        username = flask_session.get("username")

        if not username:
            # Let the view handle unauthenticated users
            return view_func(*args, **kwargs)

        try:
            # Try to get or create session
            if db_manager.is_user_connected(username):
                g.db_session = db_manager.get_session(username)
            else:
                # Database not open - for encrypted DBs this means user needs to re-login
                if db_manager.has_encryption:
                    # Clear session to force re-login
                    flask_session.clear()
                    from flask import redirect, url_for

                    return redirect(url_for("auth.login"))
                # Try to reopen unencrypted database
                logger.warning(
                    f"Reopening unencrypted database for {username} - "
                    "ensure this is intentional"
                )
                engine = db_manager.open_user_database(
                    username, UNENCRYPTED_DB_PLACEHOLDER
                )
                if engine:
                    g.db_session = db_manager.get_session(username)

        except Exception:
            logger.exception(
                f"Failed to ensure database session for {username}"
            )
            return jsonify({"error": "Database session unavailable"}), 500

        return view_func(*args, **kwargs)

    return wrapper


class DatabaseAccessMixin:
    """
    Mixin class for services that need database access.
    Provides convenient methods for database operations.
    """

    def get_db_session(
        self, username: Optional[str] = None
    ) -> Optional[Session]:
        """
        DEPRECATED: This method returns a closed session due to context manager exit.

        Use `with get_user_db_session(username) as session:` instead.

        Raises:
            DeprecationWarning: Always raised to prevent usage of broken method.
        """
        raise DeprecationWarning(
            "get_db_session() is deprecated and returns a closed session. "
            "Use `with get_user_db_session(username) as session:` instead."
        )

    @with_user_database
    def execute_with_db(
        self, db_session: Session, query_func: Callable, *args, **kwargs
    ):
        """Execute a function with database session."""
        return query_func(db_session, *args, **kwargs)
