"""Context propagation patches for OpenInference instrumentors.

Fixes two issues when multiple instrumentors coexist:

1. LangChain + OpenAI: OpenAI instrumentor's ChatCompletion span should be a
   child of LangChain's ChatOpenAI span, but LangChain uses callbacks (not
   ContextVar) for span management. Fix: patch OpenAI instrumentor to query
   LangChain's get_current_span() before creating its own span.

2. CrewAI trace isolation: CrewAI's event bus fires events in a background
   ThreadPoolExecutor thread. That thread has no OTel context from the main
   asyncio task. Fix: patch CrewAI's EventAssembler to remember the ambient
   context at construction time and use it as fallback when the background
   thread's context is empty.

These patches are applied AFTER instrumentors are activated (since we need to
patch the already-instantiated objects) and are idempotent.
"""
from __future__ import annotations

import logging
from typing import Optional

logger = logging.getLogger(__name__)


def apply_context_patches() -> None:
    """Apply all context propagation patches. Safe to call multiple times."""
    _patch_openai_langchain_bridge()
    _patch_crewai_ambient_context()


def _patch_openai_langchain_bridge() -> None:
    """Patch OpenAI instrumentor to use LangChain's active span as parent.

    Without this, ChatCompletion spans from openai instrumentor land under
    the root span instead of under LangChain's ChatOpenAI span.

    Strategy: wrap _Request.__call__ and _AsyncRequest.__call__ to attach
    the LangChain parent context before the original method runs (which
    calls _start_as_current_span internally). This way the tracer.start_span()
    inside naturally picks up the correct parent from the ambient context.
    """
    try:
        from openinference.instrumentation.openai import _request as openai_req
    except ImportError:
        return  # openai instrumentor not installed

    # Check if already patched
    if getattr(openai_req, "_eo_langchain_bridge_patched", False):
        return

    try:
        from opentelemetry import trace as trace_api
        from opentelemetry import context as context_api
        from opentelemetry.trace import INVALID_SPAN

        def _get_langchain_parent_context() -> Optional[context_api.Context]:
            """Query LangChain instrumentor for the current active span."""
            try:
                from openinference.instrumentation.langchain import get_current_span
                parent_span = get_current_span()
                if parent_span is not None and parent_span is not INVALID_SPAN:
                    return trace_api.set_span_in_context(parent_span)
            except ImportError:
                pass
            except Exception:
                pass
            return None

        # Patch _Request.__call__: attach LangChain context before execution
        _orig_request_call = openai_req._Request.__call__

        def _patched_request_call(self, wrapped, instance, args, kwargs):
            parent_ctx = _get_langchain_parent_context()
            if parent_ctx is None:
                return _orig_request_call(self, wrapped, instance, args, kwargs)
            # Attach LangChain's span as ambient context so start_span picks it up
            token = context_api.attach(parent_ctx)
            try:
                return _orig_request_call(self, wrapped, instance, args, kwargs)
            finally:
                context_api.detach(token)

        # Patch _AsyncRequest.__call__
        _orig_async_request_call = openai_req._AsyncRequest.__call__

        async def _patched_async_request_call(self, wrapped, instance, args, kwargs):
            parent_ctx = _get_langchain_parent_context()
            if parent_ctx is None:
                return await _orig_async_request_call(self, wrapped, instance, args, kwargs)
            token = context_api.attach(parent_ctx)
            try:
                return await _orig_async_request_call(self, wrapped, instance, args, kwargs)
            finally:
                context_api.detach(token)

        openai_req._Request.__call__ = _patched_request_call
        openai_req._AsyncRequest.__call__ = _patched_async_request_call
        openai_req._eo_langchain_bridge_patched = True
        logger.debug("[observability] patched openai instrumentor with langchain bridge")
    except Exception as e:
        logger.debug(f"[observability] langchain bridge patch failed: {e}")


def _patch_crewai_ambient_context() -> None:
    """Patch CrewAI's kickoff wrapper to inherit the request's OTel context.

    Problem: CrewAI's crew.kickoff() runs in run_in_executor (thread pool).
    Python's run_in_executor does NOT propagate OTel context to worker threads.
    So the CrewAI instrumentor's wrapper calls start_as_current_span() but
    gets an empty context — creating a new trace instead of joining the request's.

    Fix: patch _CrewKickoffWrapper.__call__ to attach our stored request
    context before the original logic runs.
    """
    try:
        from openinference.instrumentation.crewai import _wrappers as wrappers_mod
    except ImportError:
        return  # crewai instrumentor not installed

    if getattr(wrappers_mod, "_eo_ambient_context_patched", False):
        return

    try:
        from opentelemetry import context as context_api

        _CrewKickoffWrapper = getattr(wrappers_mod, "_CrewKickoffWrapper", None)
        if _CrewKickoffWrapper is None:
            return

        _orig_call = _CrewKickoffWrapper.__call__

        def _patched_kickoff_call(self, wrapped, instance, args, kwargs):
            from . import context_patches as _cp
            captured = getattr(_cp, "_crewai_request_context", None)
            if captured is not None:
                token = context_api.attach(captured)
                try:
                    return _orig_call(self, wrapped, instance, args, kwargs)
                finally:
                    context_api.detach(token)
            return _orig_call(self, wrapped, instance, args, kwargs)

        _CrewKickoffWrapper.__call__ = _patched_kickoff_call
        wrappers_mod._eo_ambient_context_patched = True
        logger.debug("[observability] patched crewai kickoff wrapper with ambient context")
    except Exception as e:
        logger.debug(f"[observability] crewai ambient context patch failed: {e}")


# Module-level variable for adapter to set per-request context.
# The adapter sets this AFTER attaching the root span context,
# so CrewAI's background thread can read it.
_crewai_request_context: Optional["context_api.Context"] = None


def set_request_context(ctx) -> None:
    """Called by adapter after attaching root span context."""
    global _crewai_request_context
    _crewai_request_context = ctx


def clear_request_context() -> None:
    """Called by adapter after request completes."""
    global _crewai_request_context
    _crewai_request_context = None