"""
Utility functions for handling thread-local context propagation.

This module provides helpers for propagating research context across thread boundaries,
which is necessary when strategies use ThreadPoolExecutor for parallel searches.
"""

import functools
from contextlib import contextmanager
from threading import local
from typing import Any, Callable, Dict, Generator

from loguru import logger

_g_thread_data = local()
"""
Thread-local storage for research context data.
"""


def set_search_context(context: Dict[str, Any]) -> None:
    """
    Sets the research context for this entire thread.

    Args:
        context: The context to set.

    """
    global _g_thread_data
    if hasattr(_g_thread_data, "context"):
        logger.debug(
            "Context already set for this thread. It will be overwritten."
        )
    _g_thread_data.context = context.copy()


def clear_search_context() -> None:
    """
    Clears the research context for this thread.

    Should be called in a finally block after set_search_context() to prevent
    context from leaking to subsequent tasks when threads are reused in a pool.
    """
    global _g_thread_data
    if hasattr(_g_thread_data, "context"):
        del _g_thread_data.context


def get_search_context() -> Dict[str, Any] | None:
    """
    Gets the current research context for this thread.

    Returns:
        The context dictionary, or None if no context is set.

    """
    context = getattr(_g_thread_data, "context", None)
    if context is not None:
        context = context.copy()
    return context


@contextmanager
def search_context(context: Dict[str, Any]) -> Generator[None, None, None]:
    """Context manager that sets and clears search context automatically.

    Ensures cleanup even if an exception occurs, preventing context leaks
    when threads are reused in a pool.

    Example:
        with search_context({"research_id": "123"}):
            results = engine.run(query)
    """
    set_search_context(context)
    try:
        yield
    finally:
        clear_search_context()


def preserve_research_context(func: Callable) -> Callable:
    """
    Decorator that preserves research context across thread boundaries.

    Use this decorator on functions that will be executed in ThreadPoolExecutor
    to ensure the research context (including research_id) is properly propagated.

    When metrics are disabled (e.g., in programmatic mode), this decorator
    safely does nothing to avoid database dependencies.

    Example:
        @preserve_research_context
        def search_task(query):
            return search_engine.run(query)
    """
    # Try to capture current context, but don't fail if it's not set. There
    # are legitimate cases where it might not be set, such as for
    # programmatic access.
    context = get_search_context()

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        if context is not None:
            set_search_context(context)

        try:
            return func(*args, **kwargs)
        finally:
            if context is not None:
                clear_search_context()
                # Clean up thread-local DB engines created by metrics recording
                try:
                    from ..database.thread_local_session import (
                        cleanup_current_thread,
                    )

                    cleanup_current_thread()
                except Exception:
                    logger.debug(
                        "preserve_research_context: error during cleanup_current_thread",
                        exc_info=True,
                    )

    return wrapper
