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

from ..metrics.search_tracker import get_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.

    Example:
        @preserve_research_context
        def search_task(query):
            return search_engine.run(query)
    """

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        # The context should already be captured in the closure when the decorator runs
        # Set it in the new thread
        tracker = get_search_tracker()
        if hasattr(wrapper, "_research_context"):
            tracker.set_research_context(wrapper._research_context)
        return func(*args, **kwargs)

    # Capture the current context when the decorator is applied
    wrapper._research_context = get_search_tracker()._get_research_context()
    return wrapper


def create_context_preserving_wrapper(
    func: Callable, context: Dict[str, Any] = None
) -> Callable:
    """
    Create a wrapper function that preserves research context.

    This is useful when you need to create the wrapper dynamically and can't use a decorator.

    Args:
        func: The function to wrap
        context: Optional explicit context to use. If None, captures current context.

    Returns:
        A wrapped function that sets the research context before executing
    """
    # Capture context at wrapper creation time if not provided
    if context is None:
        context = get_search_tracker()._get_research_context()

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        # Set the captured context in the new thread
        get_search_tracker().set_research_context(context)
        return func(*args, **kwargs)

    return wrapper


def run_with_context(
    func: Callable, *args, context: Dict[str, Any] = None, **kwargs
) -> Any:
    """
    Run a function with a specific research context.

    Args:
        func: The function to run
        *args: Positional arguments for the function
        context: Optional explicit context. If None, uses current context.
        **kwargs: Keyword arguments for the function

    Returns:
        The result of the function call
    """
    tracker = get_search_tracker()

    # Save current context
    original_context = tracker._get_research_context()

    try:
        # Set new context
        if context is None:
            context = original_context
        tracker.set_research_context(context)

        # Run the function
        return func(*args, **kwargs)
    finally:
        # Restore original context
        tracker.set_research_context(original_context)
