"""
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 threading import local
from typing import Any, Callable, Dict

from loguru import logger

# Lazy import to avoid database initialization for programmatic access
_search_tracker = None

_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.warning(
            "Context already set for this thread. It will be overwritten."
        )
    _g_thread_data.context = context.copy()


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


def _get_search_tracker_if_needed():
    """Get search tracker only if metrics are enabled."""
    global _search_tracker
    if _search_tracker is None:
        try:
            from ..metrics.search_tracker import get_search_tracker

            _search_tracker = get_search_tracker()
        except (ImportError, RuntimeError) as e:
            # If import fails due to database issues, metrics are disabled
            from loguru import logger

            logger.debug(
                f"Metrics tracking disabled - search tracker not available: {e}"
            )
            return None
    return _search_tracker


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)

        return func(*args, **kwargs)

    return wrapper
