"""
Utility functions for handling research context propagation.

This module provides helpers for propagating research context across thread
and asyncio boundaries. Built on ``contextvars.ContextVar`` so the context
is correctly inherited by frameworks that copy the context to worker
threads (e.g. langchain's ``ContextThreadPoolExecutor`` used by LangGraph
for parallel tool execution). For stdlib ``ThreadPoolExecutor`` — which
does not copy context — use ``preserve_research_context`` below.
"""

import functools
from contextlib import contextmanager
from contextvars import ContextVar
from typing import Any, Callable, Dict, Generator, Optional

from loguru import logger

_search_context_var: ContextVar[Optional[Dict[str, Any]]] = ContextVar(
    "ldr_search_context", default=None
)


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

    Args:
        context: The context to set.

    """
    if _search_context_var.get() is not None:
        logger.debug(
            "Context already set for this thread. It will be overwritten."
        )
    _search_context_var.set(context.copy())


def clear_search_context() -> None:
    """
    Clears the research context for the current execution context.

    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.
    """
    _search_context_var.set(None)


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

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

    """
    context = _search_context_var.get()
    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()

    # Capture the submitter thread's egress audit-hook context too. Unlike the
    # search context (a ContextVar that langchain's executors copy), the audit
    # context lives in a threading.local that stdlib ThreadPoolExecutor workers
    # do NOT inherit — so without re-arming it here the PEP-578 socket backstop
    # would be inactive on every pool worker. Capture once on the submitter
    # thread; re-arm + clear per task below.
    try:
        from ..security.egress.audit_hook import (
            get_active_context as _get_egress_ctx,
        )

        egress_ctx = _get_egress_ctx()
    except Exception:
        egress_ctx = None

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        if context is not None:
            set_search_context(context)
        if egress_ctx is not None:
            from ..security.egress.audit_hook import (
                set_active_context as _set_egress_ctx,
            )

            _set_egress_ctx(egress_ctx)

        try:
            return func(*args, **kwargs)
        finally:
            if egress_ctx is not None:
                from ..security.egress.audit_hook import (
                    clear_active_context as _clear_egress_ctx,
                )

                _clear_egress_ctx()
            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
