"""
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, 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


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

    pass


@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_app_context():
            username = flask_session.get("username")

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

        # First, check if we have a session in Flask's g object (best performance)
        if has_app_context() and hasattr(g, "db_session") and g.db_session:
            # Use existing session from g - no need to close
            session = g.db_session
            needs_close = False
        else:
            # Get password if not provided
            if not password and has_app_context():
                # Try to get from g
                if hasattr(g, "user_password"):
                    password = g.user_password
                    logger.debug(
                        f"Got password from g.user_password for {username}"
                    )
                # Try session password store
                elif flask_session.get("session_id"):
                    session_id = flask_session.get("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:
                password = "dummy"  # For unencrypted databases

            # 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:
                pass


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 username in db_manager.connections:
                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"))
                else:
                    # Try to reopen unencrypted database
                    engine = db_manager.open_user_database(username, "dummy")
                    if engine:
                        g.db_session = db_manager.get_session(username)

        except Exception:
            logger.exception(
                f"Failed to ensure database session for {username}"
            )

        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]:
        """Get database session for user."""
        try:
            with get_user_db_session(username) as session:
                return session
        except DatabaseSessionError:
            return None

    @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)
