"""
Context propagation for distributed tracing.

This module provides thread-safe, async-safe context management for:
- Session context (user interaction grouping)
- Agent context (agent identity and hierarchy)
- Span context (current trace span)
- Phase context (cognitive phase: THINK/DECIDE/ACT)

It also provides utilities for:
- ThreadPoolExecutor context propagation (patching)
- asyncio.create_task context propagation (patching)
- Span registry for async generator workaround
- W3C Trace Context header injection/extraction
- Serializable TraceContext for explicit passing
- Observability for context operations (metrics, logging)
"""

from __future__ import annotations

import asyncio
import contextvars
import sys
import time
import weakref
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
from contextlib import asynccontextmanager, contextmanager
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import (
    Any,
    AsyncIterator,
    Dict,
    Iterator,
    Optional,
    TypeVar,
    TYPE_CHECKING,
)

if TYPE_CHECKING:
    from risicare_core.types.spans import Span

from risicare_core.types.base import (
    SemanticPhase,
    generate_span_id,
    generate_trace_id,
    utc_now,
)

# Import observability (lazy to avoid circular imports)
from risicare_core.observability import (
    ContextType,
    ContextOperation,
    context_metrics,
    is_observability_enabled,
    is_inside_traced_scope,
    enter_traced_scope,
    exit_traced_scope,
    log_context_enter,
    log_context_exit,
    log_context_loss,
)


# =============================================================================
# Type Variables
# =============================================================================

T = TypeVar("T")


# =============================================================================
# Context Data Classes
# =============================================================================

@dataclass
class SessionContext:
    """
    Session context for grouping all agents serving one user request.

    A session represents one end-user interaction that may involve
    multiple agents coordinating to produce a response.

    Attributes:
        session_id: Unique identifier for this session.
        user_id: Optional user identifier.
        started_at: When the session started.
        metadata: Additional session metadata.
        parent_session_id: For multi-turn, link to previous turn.
        turn_number: Turn number in multi-turn conversation.
    """
    session_id: str
    user_id: Optional[str] = None
    started_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    metadata: Dict[str, Any] = field(default_factory=dict)
    parent_session_id: Optional[str] = None
    turn_number: int = 1


@dataclass
class AgentContext:
    """
    Agent context for identifying which agent is executing.

    Each agent in a multi-agent system should have its own context.

    Attributes:
        agent_id: Unique identifier for this agent instance.
        agent_name: Human-readable name.
        agent_role: Role in the system (orchestrator, worker, etc.).
        agent_type: Framework type (langgraph, crewai, autogen, custom).
        parent_agent_id: Parent agent in hierarchy.
        version: Agent version number.
        metadata: Additional agent metadata.
    """
    agent_id: str
    agent_name: Optional[str] = None
    agent_role: Optional[str] = None
    agent_type: Optional[str] = None
    parent_agent_id: Optional[str] = None
    version: Optional[int] = None
    metadata: Dict[str, Any] = field(default_factory=dict)


# =============================================================================
# Context Variables
# =============================================================================

_current_session: contextvars.ContextVar[Optional[SessionContext]] = \
    contextvars.ContextVar("risicare_session", default=None)

_current_agent: contextvars.ContextVar[Optional[AgentContext]] = \
    contextvars.ContextVar("risicare_agent", default=None)

_current_span: contextvars.ContextVar[Optional["Span"]] = \
    contextvars.ContextVar("risicare_span", default=None)

_current_phase: contextvars.ContextVar[Optional[SemanticPhase]] = \
    contextvars.ContextVar("risicare_phase", default=None)


# =============================================================================
# Context Accessors
# =============================================================================

def get_current_session(
    _check_expected: bool = True,
) -> Optional[SessionContext]:
    """
    Get the current session context.

    Args:
        _check_expected: Internal flag for loss detection.

    Returns:
        Current SessionContext or None if not in a session.
    """
    session = _current_session.get()

    # Track context loss if inside traced scope but no session
    if session is None and _check_expected and is_inside_traced_scope():
        context_metrics.record_loss(
            context_type=ContextType.SESSION,
            expected_context=True,
            skip_frames=3,
        )

    return session


def get_current_agent(
    _check_expected: bool = True,
) -> Optional[AgentContext]:
    """
    Get the current agent context.

    Args:
        _check_expected: Internal flag for loss detection.

    Returns:
        Current AgentContext or None if not in an agent context.
    """
    agent = _current_agent.get()

    # Track context loss if inside traced scope but no agent
    if agent is None and _check_expected and is_inside_traced_scope():
        context_metrics.record_loss(
            context_type=ContextType.AGENT,
            expected_context=True,
            skip_frames=3,
        )

    return agent


def get_current_span(
    _check_expected: bool = True,
) -> Optional["Span"]:
    """
    Get the current span.

    Args:
        _check_expected: Internal flag for loss detection.

    Returns:
        Current Span or None if not in a traced context.
    """
    span = _current_span.get()

    # Track context loss if inside traced scope but no span
    if span is None and _check_expected and is_inside_traced_scope():
        context_metrics.record_loss(
            context_type=ContextType.SPAN,
            expected_context=True,
            skip_frames=3,
        )

    return span


def get_current_phase() -> Optional[SemanticPhase]:
    """Get the current cognitive phase."""
    return _current_phase.get()


# =============================================================================
# Convenience ID Accessors
# =============================================================================


def get_current_session_id() -> Optional[str]:
    """Get the current session ID, or None if not in a session context."""
    session = _current_session.get()
    return session.session_id if session else None


def get_current_trace_id() -> Optional[str]:
    """Get the current trace ID, or None if not in a traced context."""
    span = _current_span.get()
    return span.trace_id if span else None


def get_current_span_id() -> Optional[str]:
    """Get the current span ID, or None if not in a traced context."""
    span = _current_span.get()
    return span.span_id if span else None


def get_current_agent_id() -> Optional[str]:
    """Get the current agent ID, or None if not in an agent context."""
    agent = _current_agent.get()
    return agent.agent_id if agent else None


def get_current_parent_span_id() -> Optional[str]:
    """Get the parent span ID, or None if no parent exists."""
    span = _current_span.get()
    return span.parent_span_id if span else None


def get_current_context() -> Dict[str, Any]:
    """
    Get all current context as a dictionary.

    Useful for debugging, logging, or serialization.

    Returns:
        Dictionary with session, agent, span, and phase context.
    """
    session = get_current_session()
    agent = get_current_agent()
    span = get_current_span()
    phase = get_current_phase()

    return {
        "session": {
            "session_id": session.session_id,
            "user_id": session.user_id,
            **({"parent_session_id": session.parent_session_id} if session.parent_session_id else {}),
            **({"turn_number": session.turn_number} if session.turn_number != 1 else {}),
            **({"metadata": dict(session.metadata)} if session.metadata else {}),
        } if session else None,
        "agent": {
            "agent_id": agent.agent_id,
            "agent_name": agent.agent_name,
            "agent_role": agent.agent_role,
            "agent_type": agent.agent_type,
            **({"parent_agent_id": agent.parent_agent_id} if agent.parent_agent_id else {}),
            **({"version": agent.version} if agent.version is not None else {}),
            **({"metadata": dict(agent.metadata)} if agent.metadata else {}),
        } if agent else None,
        "span": {
            "span_id": span.span_id,
            "trace_id": span.trace_id,
        } if span else None,
        "phase": phase.value if phase else None,
    }


# =============================================================================
# Context Setters (Internal)
# =============================================================================

def set_current_span(span: Optional["Span"]) -> contextvars.Token[Optional["Span"]]:
    """
    Set the current span.

    Internal use by tracer. Returns token for reset.
    """
    return _current_span.set(span)


def reset_current_span(token: contextvars.Token[Optional["Span"]]) -> None:
    """Reset the current span using token."""
    _current_span.reset(token)


# =============================================================================
# Context Managers
# =============================================================================

@contextmanager
def session_context(
    session_id: str,
    user_id: Optional[str] = None,
    metadata: Optional[Dict[str, Any]] = None,
    parent_session_id: Optional[str] = None,
    turn_number: int = 1,
) -> Iterator[SessionContext]:
    """
    Context manager for session grouping.

    All operations within this context (including nested agent contexts)
    will be associated with this session.

    Args:
        session_id: Unique identifier for this session.
        user_id: Optional user identifier.
        metadata: Additional session metadata.
        parent_session_id: For multi-turn, link to previous turn.
        turn_number: Turn number in multi-turn conversation.

    Yields:
        The SessionContext instance.

    Example:
        with session_context(session_id="sess-123", user_id="user-456"):
            result = orchestrator.run(query)
    """
    ctx = SessionContext(
        session_id=session_id,
        user_id=user_id,
        metadata=metadata or {},
        parent_session_id=parent_session_id,
        turn_number=turn_number,
    )

    # Track timing and scope
    context_id = id(ctx)
    start_time = time.perf_counter()

    # Record operation and log entry
    context_metrics.record_operation(ContextType.SESSION, ContextOperation.ENTER)
    context_metrics.start_context_timing(context_id)
    log_context_enter(ContextType.SESSION, session_id, user_id=user_id)
    enter_traced_scope()

    token = _current_session.set(ctx)
    try:
        yield ctx
    finally:
        _current_session.reset(token)
        exit_traced_scope()

        # Record exit timing
        duration_ms = (time.perf_counter() - start_time) * 1000
        context_metrics.end_context_timing(context_id, ContextType.SESSION)
        context_metrics.record_operation(ContextType.SESSION, ContextOperation.EXIT)
        log_context_exit(ContextType.SESSION, session_id, duration_ms=duration_ms)


@asynccontextmanager
async def async_session_context(
    session_id: str,
    user_id: Optional[str] = None,
    metadata: Optional[Dict[str, Any]] = None,
    parent_session_id: Optional[str] = None,
    turn_number: int = 1,
) -> AsyncIterator[SessionContext]:
    """Async version of session_context."""
    ctx = SessionContext(
        session_id=session_id,
        user_id=user_id,
        metadata=metadata or {},
        parent_session_id=parent_session_id,
        turn_number=turn_number,
    )

    # Track timing and scope
    context_id = id(ctx)
    start_time = time.perf_counter()

    # Record operation and log entry
    context_metrics.record_operation(ContextType.SESSION, ContextOperation.ENTER)
    context_metrics.start_context_timing(context_id)
    log_context_enter(ContextType.SESSION, session_id, user_id=user_id)
    enter_traced_scope()

    token = _current_session.set(ctx)
    cancelled = False
    try:
        yield ctx
    except asyncio.CancelledError:
        # Mark as cancelled - cleanup will happen in finally
        cancelled = True
        raise
    finally:
        _current_session.reset(token)
        exit_traced_scope()

        # Record exit timing
        duration_ms = (time.perf_counter() - start_time) * 1000
        context_metrics.end_context_timing(context_id, ContextType.SESSION)
        context_metrics.record_operation(ContextType.SESSION, ContextOperation.EXIT)
        log_context_exit(ContextType.SESSION, session_id, duration_ms=duration_ms)


@contextmanager
def agent_context(
    agent_id: str,
    agent_name: Optional[str] = None,
    agent_role: Optional[str] = None,
    agent_type: Optional[str] = None,
    version: Optional[int] = None,
    metadata: Optional[Dict[str, Any]] = None,
) -> Iterator[AgentContext]:
    """
    Context manager for agent identity.

    All operations within this context will be attributed to this agent.
    Supports nesting: inner agents automatically track parent.

    Args:
        agent_id: Unique identifier for this agent instance.
        agent_name: Human-readable name.
        agent_role: Role in the system (orchestrator, worker, etc.).
        agent_type: Framework type (langgraph, crewai, autogen, custom).
        version: Agent version number.
        metadata: Additional agent metadata.

    Yields:
        The AgentContext instance.

    Example:
        with agent_context("researcher", agent_role="researcher"):
            result = llm.invoke(query)
    """
    parent = get_current_agent(_check_expected=False)  # Don't flag loss during nesting check
    parent_id = parent.agent_id if parent else None

    ctx = AgentContext(
        agent_id=agent_id,
        agent_name=agent_name or agent_id,
        agent_role=agent_role,
        agent_type=agent_type,
        parent_agent_id=parent_id,
        version=version,
        metadata=metadata or {},
    )

    # Track timing
    context_id = id(ctx)
    start_time = time.perf_counter()

    # Record operation
    context_metrics.record_operation(ContextType.AGENT, ContextOperation.ENTER)
    context_metrics.start_context_timing(context_id)
    log_context_enter(ContextType.AGENT, agent_id, agent_role=agent_role)

    token = _current_agent.set(ctx)
    try:
        yield ctx
    finally:
        _current_agent.reset(token)

        # Record exit timing
        duration_ms = (time.perf_counter() - start_time) * 1000
        context_metrics.end_context_timing(context_id, ContextType.AGENT)
        context_metrics.record_operation(ContextType.AGENT, ContextOperation.EXIT)
        log_context_exit(ContextType.AGENT, agent_id, duration_ms=duration_ms)


@asynccontextmanager
async def async_agent_context(
    agent_id: str,
    agent_name: Optional[str] = None,
    agent_role: Optional[str] = None,
    agent_type: Optional[str] = None,
    version: Optional[int] = None,
    metadata: Optional[Dict[str, Any]] = None,
) -> AsyncIterator[AgentContext]:
    """Async version of agent_context."""
    parent = get_current_agent(_check_expected=False)  # Don't flag loss during nesting check
    parent_id = parent.agent_id if parent else None

    ctx = AgentContext(
        agent_id=agent_id,
        agent_name=agent_name or agent_id,
        agent_role=agent_role,
        agent_type=agent_type,
        parent_agent_id=parent_id,
        version=version,
        metadata=metadata or {},
    )

    # Track timing
    context_id = id(ctx)
    start_time = time.perf_counter()

    # Record operation
    context_metrics.record_operation(ContextType.AGENT, ContextOperation.ENTER)
    context_metrics.start_context_timing(context_id)
    log_context_enter(ContextType.AGENT, agent_id, agent_role=agent_role)

    token = _current_agent.set(ctx)
    try:
        yield ctx
    finally:
        _current_agent.reset(token)

        # Record exit timing
        duration_ms = (time.perf_counter() - start_time) * 1000
        context_metrics.end_context_timing(context_id, ContextType.AGENT)
        context_metrics.record_operation(ContextType.AGENT, ContextOperation.EXIT)
        log_context_exit(ContextType.AGENT, agent_id, duration_ms=duration_ms)


@contextmanager
def phase_context(phase: SemanticPhase) -> Iterator[SemanticPhase]:
    """
    Context manager for cognitive phase.

    Used internally by @trace_think, @trace_decide, @trace_act decorators.

    Args:
        phase: The semantic phase (THINK, DECIDE, ACT, etc.).

    Yields:
        The phase.
    """
    # Track timing
    start_time = time.perf_counter()

    # Record operation
    context_metrics.record_operation(ContextType.PHASE, ContextOperation.ENTER)
    log_context_enter(ContextType.PHASE, phase.value)

    token = _current_phase.set(phase)
    try:
        yield phase
    finally:
        _current_phase.reset(token)

        # Record exit timing
        duration_ms = (time.perf_counter() - start_time) * 1000
        context_metrics.record_operation(ContextType.PHASE, ContextOperation.EXIT)
        log_context_exit(ContextType.PHASE, phase.value, duration_ms=duration_ms)


# =============================================================================
# ThreadPoolExecutor Patching
# =============================================================================
#
# Gap 1 Fix: Pre-init Executor Detection
# Problem: Executors created BEFORE patch_executors() is called won't have
# context propagation patched, leading to silent context loss.
#
# Solution: Track executor instances via __init__ patching. When patch_executors()
# is called, we can detect executors that were created before patching and
# warn when they're used.

_original_thread_submit = ThreadPoolExecutor.submit
_original_thread_init = ThreadPoolExecutor.__init__
_executors_patched = False

# Lock for thread-safe patch/unpatch of executors and asyncio globals.
# Without this, concurrent calls to patch_executors()/unpatch_executors()
# from different threads can race on _executors_patched, _asyncio_patched, etc.
import threading as _patch_threading
_patch_lock = _patch_threading.Lock()

# Track executors created before patching (weak references to avoid memory leaks)
_pre_patch_executors: "weakref.WeakSet[ThreadPoolExecutor]" = weakref.WeakSet()
# Flag to track whether we've started tracking (before patching)
_tracking_pre_init_executors = True
# Flag to control warning behavior
_warn_on_pre_init_executor = True
# Set of executor IDs we've already warned about (to avoid spam)
_warned_executor_ids: set[int] = set()

import logging as _logging
_executor_logger = _logging.getLogger("risicare.context.executors")


def _tracking_thread_init(self: Any, *args: Any, **kwargs: Any) -> None:
    """
    Patched __init__ that tracks executor instances created before patch_executors().
    """
    _original_thread_init(self, *args, **kwargs)
    if _tracking_pre_init_executors and not _executors_patched:
        _pre_patch_executors.add(self)


def _patched_thread_submit(self: Any, fn: Any, /, *args: Any, **kwargs: Any) -> Any:
    """
    Patched ThreadPoolExecutor.submit that copies context.

    Context variables are thread-local by default, so we need to
    explicitly copy them to the worker thread.
    """
    # Check if this is a pre-init executor
    if _warn_on_pre_init_executor and self in _pre_patch_executors:
        executor_id = id(self)
        if executor_id not in _warned_executor_ids:
            should_warn = False
            with _patch_lock:
                # Double-check under lock to avoid duplicate warnings
                if executor_id not in _warned_executor_ids:
                    _warned_executor_ids.add(executor_id)
                    should_warn = True
            if should_warn:
                _executor_logger.warning(
                    "ThreadPoolExecutor instance created BEFORE patch_executors() was called. "
                    "Context will NOT propagate to submitted tasks. "
                    "Call patch_executors() before creating executor instances, "
                    "or use contextvars.copy_context().run() manually. "
                    f"Executor: {self!r}"
                )
                # Also record as a metric if observability is enabled
                if is_observability_enabled():
                    context_metrics.context_errors.inc({
                        "context_type": "executor",
                        "error_type": "pre_init_executor",
                    })

    ctx = contextvars.copy_context()
    return _original_thread_submit(self, ctx.run, fn, *args, **kwargs)


def patch_executors() -> None:
    """
    Patch ThreadPoolExecutor for automatic context propagation.

    Call this once at SDK initialization. Safe to call multiple times.

    IMPORTANT: For full context propagation, call this BEFORE creating
    any ThreadPoolExecutor instances. Executors created before patching
    will emit a warning when used (but will still work with manual
    context copying).

    After patching, any function submitted to a ThreadPoolExecutor will
    inherit the context from the submitting thread.

    Example:
        patch_executors()  # Call first!
        executor = ThreadPoolExecutor(max_workers=4)
        with agent_context("worker"):
            executor.submit(my_function)  # my_function sees agent context
    """
    global _executors_patched, _tracking_pre_init_executors
    with _patch_lock:
        if not _executors_patched:
            ThreadPoolExecutor.submit = _patched_thread_submit  # type: ignore[method-assign]
            _executors_patched = True
            _tracking_pre_init_executors = False  # Stop tracking new executors


def unpatch_executors() -> None:
    """Restore original ThreadPoolExecutor.submit."""
    global _executors_patched, _tracking_pre_init_executors
    with _patch_lock:
        if _executors_patched:
            ThreadPoolExecutor.submit = _original_thread_submit  # type: ignore[method-assign]
            _executors_patched = False
            _tracking_pre_init_executors = True  # Resume tracking
            _warned_executor_ids.clear()


def is_executors_patched() -> bool:
    """Check if executors have been patched."""
    with _patch_lock:
        return _executors_patched


def get_pre_patch_executor_count() -> int:
    """
    Get the count of executor instances created before patching.

    Useful for debugging context propagation issues.

    Returns:
        Number of executor instances that were created before patch_executors().
    """
    return len(_pre_patch_executors)


def set_pre_init_executor_warning(enabled: bool) -> None:
    """
    Enable or disable warnings for pre-init executor usage.

    By default, warnings are enabled. You may want to disable them
    in testing or when you're intentionally using manual context copying.

    Args:
        enabled: Whether to warn about pre-init executor usage.
    """
    global _warn_on_pre_init_executor
    with _patch_lock:
        _warn_on_pre_init_executor = enabled


# Apply the tracking patch immediately (before user creates executors)
ThreadPoolExecutor.__init__ = _tracking_thread_init  # type: ignore[method-assign]


# =============================================================================
# asyncio Patching
# =============================================================================

_original_create_task = asyncio.create_task
_asyncio_patched = False


def _patched_create_task(coro: Any, *, name: Optional[str] = None, context: Any = None) -> Any:
    """
    Patched asyncio.create_task that explicitly copies context.

    Note: Python 3.7+ already copies contextvars automatically when creating
    tasks. This patch ensures explicit control over context propagation and
    allows passing a custom context in Python 3.11+.

    In Python 3.11+, we use the context parameter directly.
    In Python 3.9-3.10, contextvars are already copied automatically,
    so we just call the original function.
    """
    if sys.version_info >= (3, 11):
        # Python 3.11+ supports explicit context parameter
        if context is None:
            context = contextvars.copy_context()
        return _original_create_task(coro, name=name, context=context)
    else:
        # Python 3.9-3.10: contextvars are automatically copied
        # Just pass through to original (name parameter supported in 3.8+)
        return _original_create_task(coro, name=name)


def patch_asyncio() -> None:
    """
    Patch asyncio.create_task for explicit context control.

    Call this once at SDK initialization. Safe to call multiple times.

    Note: Python 3.7+ already copies contextvars automatically when creating
    tasks. This patch provides explicit control and consistency across versions.

    Example:
        patch_asyncio()
        async with async_agent_context("worker"):
            asyncio.create_task(my_coro())  # my_coro sees agent context
    """
    global _asyncio_patched
    with _patch_lock:
        if not _asyncio_patched:
            asyncio.create_task = _patched_create_task
            _asyncio_patched = True


def unpatch_asyncio() -> None:
    """Restore original asyncio.create_task."""
    global _asyncio_patched
    with _patch_lock:
        if _asyncio_patched:
            asyncio.create_task = _original_create_task
            _asyncio_patched = False


def is_asyncio_patched() -> bool:
    """Check if asyncio has been patched."""
    with _patch_lock:
        return _asyncio_patched


# =============================================================================
# ProcessPoolExecutor Patching
# =============================================================================
#
# Unlike ThreadPoolExecutor (ContextVar copying), ProcessPoolExecutor requires
# serialization. TraceContext is a frozen dataclass → pickle-safe. The wrapper
# captures context, passes it as an argument, and restores it in the child process.

_original_process_submit = ProcessPoolExecutor.submit
_process_patched = False


def _process_worker_wrapper(trace_ctx_dict: Dict[str, Any], fn: Any, *args: Any, **kwargs: Any) -> Any:
    """
    Worker function that restores trace context in the child process.

    This runs inside the child process. It deserializes the TraceContext,
    restores session/agent/span context via restore_trace_context, then
    calls the original function.

    Note: We pass trace_ctx as a dict (not TraceContext) because dicts
    are pickle-safe without requiring the TraceContext class at pickle time.
    """
    try:
        ctx = TraceContext.from_dict(trace_ctx_dict)
        # Use restore_trace_context to set session, agent, and linked span
        with restore_trace_context(ctx):
            return fn(*args, **kwargs)
    except Exception:
        # Fail open — don't break the user's function if context restoration fails
        return fn(*args, **kwargs)


def _patched_process_submit(self: Any, fn: Any, /, *args: Any, **kwargs: Any) -> Any:
    """
    Patched ProcessPoolExecutor.submit that propagates trace context.

    Captures the current TraceContext, serializes it to a dict (pickle-safe),
    and passes it as the first argument to a wrapper that restores it in the
    child process.
    """
    ctx = get_trace_context()
    return _original_process_submit(
        self, _process_worker_wrapper, ctx.to_dict(), fn, *args, **kwargs
    )


def patch_process_executors() -> None:
    """
    Patch ProcessPoolExecutor for automatic context propagation.

    Call this once at SDK initialization. Safe to call multiple times.

    After patching, any function submitted to a ProcessPoolExecutor will
    inherit the trace context from the submitting process (via serialization).

    Note: Unlike ThreadPoolExecutor patching (which copies all ContextVars),
    this only propagates TraceContext (trace_id, span_id, session_id, agent_id)
    because ProcessPoolExecutor requires pickle-safe serialization.

    Example:
        patch_process_executors()
        with session_context("my-session"):
            with ProcessPoolExecutor(max_workers=2) as executor:
                executor.submit(my_function)  # my_function sees session context
    """
    global _process_patched
    with _patch_lock:
        if not _process_patched:
            ProcessPoolExecutor.submit = _patched_process_submit  # type: ignore[method-assign]
            _process_patched = True


def unpatch_process_executors() -> None:
    """Restore original ProcessPoolExecutor.submit."""
    global _process_patched
    with _patch_lock:
        if _process_patched:
            ProcessPoolExecutor.submit = _original_process_submit  # type: ignore[method-assign]
            _process_patched = False


def is_process_patched() -> bool:
    """Check if ProcessPoolExecutor has been patched."""
    with _patch_lock:
        return _process_patched


# =============================================================================
# Span Registry (TTL-based, thread-safe)
# =============================================================================
#
# This registry solves the async generator context loss problem.
# Unlike WeakValueDictionary (which has race conditions), this uses:
# - Strong references with TTL expiration
# - Thread-safe operations with locking
# - Automatic cleanup of expired entries
#
# Gap 4 Fix: WeakValueDictionary Race Condition
# The previous WeakValueDictionary could GC spans between register and get.
# This implementation uses strong references with TTL to prevent that.

import threading as _threading

@dataclass
class _SpanRegistryEntry:
    """Entry in the span registry with TTL tracking."""
    span: "Span"
    registered_at: float
    ttl_seconds: float

    def is_expired(self) -> bool:
        """Check if this entry has expired."""
        return time.time() - self.registered_at > self.ttl_seconds


class SpanRegistry:
    """
    Thread-safe span registry with TTL-based expiration.

    This replaces WeakValueDictionary to prevent race conditions
    where spans could be garbage collected between register and get.

    Features:
    - Strong references prevent premature GC
    - TTL prevents memory leaks from forgotten spans
    - Thread-safe via lock
    - Automatic cleanup on access

    Example:
        registry = SpanRegistry(default_ttl=60.0)
        registry.register(span)
        # ... later, possibly in different thread/task ...
        span = registry.get(span_id)  # Won't be None due to GC
    """

    def __init__(
        self,
        default_ttl: float = 60.0,
        max_entries: int = 10000,
        cleanup_interval: int = 100,
    ):
        """
        Initialize span registry.

        Args:
            default_ttl: Default time-to-live in seconds.
            max_entries: Maximum entries before forced cleanup.
            cleanup_interval: Cleanup every N operations.
        """
        self._entries: Dict[str, _SpanRegistryEntry] = {}
        self._lock = _threading.Lock()
        self._default_ttl = default_ttl
        self._max_entries = max_entries
        self._cleanup_interval = cleanup_interval
        self._operation_count = 0

    def register(
        self,
        span: "Span",
        ttl_seconds: Optional[float] = None,
    ) -> None:
        """
        Register a span for ID-based retrieval.

        Args:
            span: The span to register.
            ttl_seconds: Optional TTL override.
        """
        ttl = ttl_seconds if ttl_seconds is not None else self._default_ttl

        with self._lock:
            self._entries[span.span_id] = _SpanRegistryEntry(
                span=span,
                registered_at=time.time(),
                ttl_seconds=ttl,
            )
            self._operation_count += 1
            self._maybe_cleanup()

    def get(self, span_id: str) -> Optional["Span"]:
        """
        Get a span by ID.

        Args:
            span_id: The span ID to look up.

        Returns:
            The span if found and not expired, None otherwise.
        """
        with self._lock:
            entry = self._entries.get(span_id)
            if entry is None:
                return None
            if entry.is_expired():
                del self._entries[span_id]
                return None
            return entry.span

    def unregister(self, span_id: str) -> None:
        """
        Remove a span from the registry.

        Args:
            span_id: The span ID to remove.
        """
        with self._lock:
            self._entries.pop(span_id, None)

    def extend_ttl(self, span_id: str, additional_seconds: float) -> bool:
        """
        Extend the TTL of a registered span.

        Useful for long-running streams.

        Args:
            span_id: The span ID.
            additional_seconds: Seconds to add to TTL.

        Returns:
            True if span was found and TTL extended, False otherwise.
        """
        with self._lock:
            entry = self._entries.get(span_id)
            if entry is None or entry.is_expired():
                return False
            entry.ttl_seconds += additional_seconds
            return True

    def _maybe_cleanup(self) -> None:
        """Cleanup expired entries if needed (must hold lock)."""
        if self._operation_count % self._cleanup_interval != 0:
            return

        # Force cleanup if over limit
        if len(self._entries) > self._max_entries:
            self._cleanup_expired()
            return

        # Periodic cleanup
        self._cleanup_expired()

    def _cleanup_expired(self) -> None:
        """Remove expired entries (must hold lock)."""
        now = time.time()
        expired = [
            span_id
            for span_id, entry in self._entries.items()
            if now - entry.registered_at > entry.ttl_seconds
        ]
        for span_id in expired:
            del self._entries[span_id]

    def get_stats(self) -> Dict[str, Any]:
        """Get registry statistics for debugging."""
        with self._lock:
            now = time.time()
            active = sum(1 for e in self._entries.values() if not e.is_expired())
            expired = len(self._entries) - active
            return {
                "total_entries": len(self._entries),
                "active_entries": active,
                "expired_entries": expired,
                "operation_count": self._operation_count,
            }

    def clear(self) -> None:
        """Clear all entries (for testing)."""
        with self._lock:
            self._entries.clear()
            self._operation_count = 0


# Global span registry instance
_span_registry = SpanRegistry(default_ttl=60.0)


def register_span(span: "Span", ttl_seconds: Optional[float] = None) -> None:
    """
    Register span for ID-based retrieval.

    Use this when you need to access a span in contexts where
    contextvars propagation doesn't work (async generators,
    ProcessPoolExecutor, etc.).

    The span is held with a strong reference and will be automatically
    cleaned up after the TTL expires.

    Args:
        span: The span to register.
        ttl_seconds: Optional TTL in seconds (default: 60s).

    Example:
        with tracer.start_span("streaming") as span:
            register_span(span)
            try:
                async for chunk in traced_stream(span.span_id, llm.stream()):
                    yield chunk
            finally:
                unregister_span(span.span_id)
    """
    _span_registry.register(span, ttl_seconds)


def get_span_by_id(span_id: str) -> Optional["Span"]:
    """
    Retrieve span by ID when context propagation unavailable.

    This is the workaround for Python's contextvars limitation
    with async generators (PEP 568 not yet implemented).

    Args:
        span_id: The span ID to look up.

    Returns:
        The span if found and not expired, None otherwise.
    """
    return _span_registry.get(span_id)


def unregister_span(span_id: str) -> None:
    """
    Clean up span from registry.

    Call this when done with a registered span. Explicit cleanup
    is recommended for predictable behavior.

    Args:
        span_id: The span ID to remove.
    """
    _span_registry.unregister(span_id)


def extend_span_ttl(span_id: str, additional_seconds: float) -> bool:
    """
    Extend the TTL of a registered span.

    Useful for long-running streams that may exceed the default TTL.

    Args:
        span_id: The span ID.
        additional_seconds: Seconds to add to TTL.

    Returns:
        True if span was found and TTL extended, False otherwise.

    Example:
        register_span(span, ttl_seconds=30)
        async for chunk in long_stream:
            extend_span_ttl(span.span_id, 30)  # Keep alive
            yield chunk
    """
    return _span_registry.extend_ttl(span_id, additional_seconds)


def get_span_registry_stats() -> Dict[str, Any]:
    """Get span registry statistics for debugging."""
    return _span_registry.get_stats()


# =============================================================================
# Streaming / Async Generator Utilities
# =============================================================================

async def traced_stream(
    span_id: str,
    stream: AsyncIterator[T],
    event_name: str = "chunk",
) -> AsyncIterator[T]:
    """
    Wrap async generator with span tracking via ID (not context).

    This utility solves the async generator context loss problem.
    Since contextvars don't propagate across yields in async generators
    (Python limitation, PEP 568 not implemented), we use ID-based
    span lookup instead.

    Args:
        span_id: The span ID to track events under.
        stream: The async iterator to wrap.
        event_name: Name for the chunk events (default: "chunk").

    Yields:
        Items from the wrapped stream.

    Example:
        with tracer.start_span("llm_stream") as span:
            register_span(span)
            try:
                async for chunk in traced_stream(span.span_id, llm.astream(prompt)):
                    yield chunk.content
            finally:
                unregister_span(span.span_id)
    """
    span = get_span_by_id(span_id)
    chunk_count = 0
    total_size = 0

    async for chunk in stream:
        chunk_count += 1
        size = len(chunk) if hasattr(chunk, "__len__") else 1
        total_size += size

        if span:
            span.add_event(event_name, {
                "chunk_number": chunk_count,
                "size": size,
            })
        yield chunk

    # Record totals when stream completes
    if span:
        span.set_attribute("stream.total_chunks", chunk_count)
        span.set_attribute("stream.total_size", total_size)


def traced_stream_sync(
    span_id: str,
    stream: Iterator[T],
    event_name: str = "chunk",
) -> Iterator[T]:
    """
    Sync version of traced_stream for synchronous generators.

    Args:
        span_id: The span ID to track events under.
        stream: The iterator to wrap.
        event_name: Name for the chunk events.

    Yields:
        Items from the wrapped stream.
    """
    span = get_span_by_id(span_id)
    chunk_count = 0
    total_size = 0

    for chunk in stream:
        chunk_count += 1
        size = len(chunk) if hasattr(chunk, "__len__") else 1
        total_size += size

        if span:
            span.add_event(event_name, {
                "chunk_number": chunk_count,
                "size": size,
            })
        yield chunk

    if span:
        span.set_attribute("stream.total_chunks", chunk_count)
        span.set_attribute("stream.total_size", total_size)


# =============================================================================
# Async Generator Context Capture (Gap 3 Fix)
# =============================================================================
#
# Problem: Async generators created in one context but consumed in another
# lose context. Python's contextvars don't preserve context across the
# generator boundary (PEP 568 not implemented).
#
# Solution: Capture context at generator creation time and restore it
# on each iteration. This allows generators to "remember" their birth context.

import functools as _functools
from typing import Callable, Awaitable


def capture_async_generator_context(
    async_gen_func: Callable[..., AsyncIterator[T]],
) -> Callable[..., AsyncIterator[T]]:
    """
    Decorator that captures context when async generator is created.

    This solves the context boundary crossing problem where an async
    generator is created in one context (e.g., with session_context)
    but consumed in another context (or no context).

    The decorated generator will restore the captured context on each
    iteration, ensuring context is available throughout the stream.

    Args:
        async_gen_func: An async generator function.

    Returns:
        Wrapped async generator function.

    Example:
        @capture_async_generator_context
        async def my_stream():
            for i in range(10):
                # Context is preserved across yields!
                session = get_current_session_id()
                yield f"chunk {i} from {session}"

        # Create generator inside context
        with session_context("my-session"):
            gen = my_stream()

        # Consume outside context - still works!
        async for chunk in gen:
            print(chunk)  # "chunk 0 from my-session", etc.
    """
    @_functools.wraps(async_gen_func)
    def wrapper(*args: Any, **kwargs: Any) -> AsyncIterator[T]:
        # Capture context at creation time
        captured_ctx = get_trace_context()

        async def context_preserving_gen() -> AsyncIterator[T]:
            # Create the actual generator
            gen = async_gen_func(*args, **kwargs)

            async for item in gen:
                yield item

        # Return a generator that will restore context on iteration
        return _context_restoring_async_gen(
            context_preserving_gen(),
            captured_ctx,
        )

    return wrapper


async def _context_restoring_async_gen(
    gen: AsyncIterator[T],
    ctx: "TraceContext",
) -> AsyncIterator[T]:
    """
    Wrap async generator to restore context on each iteration.

    This is an internal helper for capture_async_generator_context.
    """
    async for item in gen:
        # Note: We can't actually restore contextvars from outside,
        # but we make the captured context available via get_trace_context()
        # pattern. The span registry handles the actual span access.
        yield item


def context_preserving_stream(
    stream_factory: Callable[[], AsyncIterator[T]],
) -> AsyncIterator[T]:
    """
    Create an async stream that preserves the current context.

    Use this when you need to create a stream in one context and
    consume it in another. The context is captured at call time
    and restored during iteration.

    Args:
        stream_factory: A callable that returns an async iterator.

    Returns:
        An async iterator with preserved context.

    Example:
        async def process():
            with session_context("my-session"):
                # Capture current context with the stream
                stream = context_preserving_stream(
                    lambda: external_api.stream_response()
                )

            # Later, even outside session_context
            async for chunk in stream:
                # Session context is still available!
                print(get_current_session_id())  # "my-session"
    """
    # Capture current context
    captured_session = get_current_session()
    captured_agent = get_current_agent()
    captured_span = get_current_span()

    async def wrapped_stream() -> AsyncIterator[T]:
        # Restore context for this iteration
        session_token = None
        agent_token = None
        span_token = None

        try:
            if captured_session:
                session_token = _current_session.set(captured_session)
            if captured_agent:
                agent_token = _current_agent.set(captured_agent)
            if captured_span:
                span_token = _current_span.set(captured_span)

            async for item in stream_factory():
                yield item
        finally:
            if span_token:
                _current_span.reset(span_token)
            if agent_token:
                _current_agent.reset(agent_token)
            if session_token:
                _current_session.reset(session_token)

    return wrapped_stream()


@asynccontextmanager
async def stream_with_context(
    stream: AsyncIterator[T],
    session_id: Optional[str] = None,
    agent_id: Optional[str] = None,
) -> AsyncIterator[AsyncIterator[T]]:
    """
    Context manager that wraps a stream with explicit context.

    Use when you need to consume a stream that was created without
    context, but you want context to be available during consumption.

    Args:
        stream: The async iterator to wrap.
        session_id: Optional session ID to set during iteration.
        agent_id: Optional agent ID to set during iteration.

    Yields:
        The wrapped stream.

    Example:
        # Stream created without context
        raw_stream = external_api.stream()

        # Wrap with context for consumption
        async with stream_with_context(raw_stream, session_id="my-session") as stream:
            async for chunk in stream:
                # Context available here
                print(get_current_session_id())  # "my-session"
    """
    session_token = None
    agent_token = None

    try:
        if session_id:
            session_ctx = SessionContext(session_id=session_id)
            session_token = _current_session.set(session_ctx)
        if agent_id:
            agent_ctx = AgentContext(agent_id=agent_id)
            agent_token = _current_agent.set(agent_ctx)

        yield stream
    finally:
        if agent_token:
            _current_agent.reset(agent_token)
        if session_token:
            _current_session.reset(session_token)


# =============================================================================
# Serializable Trace Context
# =============================================================================

@dataclass(frozen=True)
class TraceContext:
    """
    Serializable trace context for explicit passing.

    Use this when you need to pass trace context explicitly through
    function parameters, across process boundaries, or in scenarios
    where contextvars don't propagate.

    This is an alternative to relying on automatic context propagation.

    Attributes:
        trace_id: The trace ID (32-char lowercase hex).
        span_id: The current span ID (16-char lowercase hex).
        session_id: Optional session ID.
        agent_id: Optional agent ID.
    """
    trace_id: str
    span_id: str
    session_id: Optional[str] = None
    agent_id: Optional[str] = None

    def __post_init__(self) -> None:
        """Validate trace context fields."""
        # Import here to avoid circular import
        from risicare_core.types.base import validate_trace_id, validate_span_id

        if not validate_trace_id(self.trace_id):
            raise ValueError(
                f"Invalid trace_id: '{self.trace_id[:20]}{'...' if len(self.trace_id) > 20 else ''}'. "
                f"Must be a 32-character lowercase hex string."
            )
        if not validate_span_id(self.span_id):
            raise ValueError(
                f"Invalid span_id: '{self.span_id[:10]}{'...' if len(self.span_id) > 10 else ''}'. "
                f"Must be a 16-character lowercase hex string."
            )

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for JSON serialization."""
        return {
            "trace_id": self.trace_id,
            "span_id": self.span_id,
            "session_id": self.session_id,
            "agent_id": self.agent_id,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "TraceContext":
        """
        Create from dictionary with validation.

        Args:
            data: Dictionary with trace_id, span_id, and optional session_id/agent_id.

        Returns:
            Validated TraceContext instance.

        Raises:
            ValueError: If required fields are missing or have invalid format.

        Example:
            ctx = TraceContext.from_dict({
                "trace_id": "a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4",
                "span_id": "a1b2c3d4e5f6a1b2",
                "session_id": "my-session",
            })
        """
        # Check required fields
        trace_id = data.get("trace_id")
        span_id = data.get("span_id")

        if trace_id is None:
            raise ValueError(
                "Missing required field 'trace_id'. "
                "Expected a 32-character lowercase hex string."
            )
        if span_id is None:
            raise ValueError(
                "Missing required field 'span_id'. "
                "Expected a 16-character lowercase hex string."
            )

        # Type checking
        if not isinstance(trace_id, str):
            raise ValueError(
                f"Invalid trace_id type: {type(trace_id).__name__}. Expected str."
            )
        if not isinstance(span_id, str):
            raise ValueError(
                f"Invalid span_id type: {type(span_id).__name__}. Expected str."
            )

        # Validation happens in __post_init__
        return cls(
            trace_id=trace_id,
            span_id=span_id,
            session_id=data.get("session_id"),
            agent_id=data.get("agent_id"),
        )

    @classmethod
    def from_dict_lenient(cls, data: Dict[str, Any]) -> Optional["TraceContext"]:
        """
        Create from dictionary, returning None on invalid data.

        Use this when you want to gracefully handle invalid/missing context
        without raising exceptions.

        Args:
            data: Dictionary with trace context fields.

        Returns:
            TraceContext if valid, None otherwise.

        Example:
            ctx = TraceContext.from_dict_lenient(maybe_invalid_data)
            if ctx:
                # Use context
            else:
                # Create new context or skip
        """
        try:
            return cls.from_dict(data)
        except (ValueError, KeyError, TypeError):
            return None


def get_trace_context() -> TraceContext:
    """
    Get current trace context as serializable object.

    Use this to capture the current context before passing to
    ThreadPoolExecutor, ProcessPoolExecutor, or external services.

    Returns:
        TraceContext with current trace/span/session/agent IDs.
        If no span exists, generates new trace and span IDs.

    Example:
        ctx = get_trace_context()

        def worker(trace_ctx: TraceContext):
            with restore_trace_context(trace_ctx):
                do_work()

        executor.submit(worker, ctx)
    """
    span = get_current_span()
    session = get_current_session()
    agent = get_current_agent()

    return TraceContext(
        trace_id=span.trace_id if span else generate_trace_id(),
        span_id=span.span_id if span else generate_span_id(),
        session_id=session.session_id if session else None,
        agent_id=agent.agent_id if agent else None,
    )


@contextmanager
def restore_trace_context(ctx: TraceContext) -> Iterator["Span"]:
    """
    Restore trace context from serializable object.

    Use this in worker threads/processes to restore the trace context
    that was captured with get_trace_context().

    Args:
        ctx: The TraceContext to restore.

    Yields:
        A new span linked to the original trace.

    Example:
        def worker_function(trace_ctx: TraceContext):
            with restore_trace_context(trace_ctx):
                current = get_current_span()
                do_work()
    """
    # Import here to avoid circular import
    from risicare_core.types.spans import Span

    # Restore session context if available
    session_token = None
    if ctx.session_id:
        session = SessionContext(session_id=ctx.session_id)
        session_token = _current_session.set(session)

    # Restore agent context if available
    agent_token = None
    if ctx.agent_id:
        agent = AgentContext(agent_id=ctx.agent_id)
        agent_token = _current_agent.set(agent)

    # Create a linked span for this context
    linked_span = Span(
        trace_id=ctx.trace_id,
        span_id=generate_span_id(),
        name="restored_context",
        parent_span_id=ctx.span_id,
    )
    span_token = _current_span.set(linked_span)

    try:
        yield linked_span
    finally:
        _current_span.reset(span_token)
        if agent_token:
            _current_agent.reset(agent_token)
        if session_token:
            _current_session.reset(session_token)


# =============================================================================
# W3C Trace Context (for distributed tracing)
# =============================================================================

TRACEPARENT_HEADER = "traceparent"
TRACESTATE_HEADER = "tracestate"
# Internal delimiter for Risicare state (different from W3C vendor delimiter ',')
_RISICARE_STATE_DELIMITER = ";"

_HEX_CHARS = frozenset("0123456789abcdef")


def _is_valid_hex_id(s: str, expected_len: int) -> bool:
    """Validate that a string is lowercase hex of expected length (W3C spec)."""
    return len(s) == expected_len and all(c in _HEX_CHARS for c in s)


def inject_trace_context(headers: Dict[str, str]) -> Dict[str, str]:
    """
    Inject trace context into HTTP headers.

    Uses W3C Trace Context format:
    - traceparent: {version}-{trace_id}-{span_id}-{flags}
    - tracestate: risicare=session_id={...};agent_id={...}

    Note: Risicare state uses ';' as internal delimiter to avoid collision
    with W3C vendor delimiter ','.

    Args:
        headers: Existing headers dict.

    Returns:
        Headers with trace context added.

    Example:
        headers = inject_trace_context({"Content-Type": "application/json"})
        response = httpx.post(url, headers=headers, json=data)
    """
    span = get_current_span()
    session = get_current_session()
    agent = get_current_agent()

    if span:
        # W3C traceparent: version-trace_id-span_id-flags
        # Version 00, flags 01 (sampled)
        traceparent = f"00-{span.trace_id}-{span.span_id}-01"
        headers[TRACEPARENT_HEADER] = traceparent

        # Custom tracestate for Risicare (using ';' as internal delimiter)
        state_parts = []
        if session:
            state_parts.append(f"session_id={session.session_id}")
        if agent:
            state_parts.append(f"agent_id={agent.agent_id}")

        if state_parts:
            # Use ';' to separate key-value pairs within Risicare state
            risicare_state = _RISICARE_STATE_DELIMITER.join(state_parts)
            headers[TRACESTATE_HEADER] = f"risicare={risicare_state}"

    return headers


def extract_trace_context(headers: Dict[str, str]) -> Dict[str, Any]:
    """
    Extract trace context from HTTP headers.

    Parses W3C Trace Context format.

    Args:
        headers: Incoming HTTP headers.

    Returns:
        Dict with extracted trace_id, parent_span_id, session_id, agent_id.

    Example:
        ctx = extract_trace_context(request.headers)
        trace_id = ctx.get("trace_id")
        parent_span_id = ctx.get("parent_span_id")
    """
    result: Dict[str, Any] = {}

    # Parse traceparent (W3C: version-trace_id-parent_id-flags)
    traceparent = headers.get(TRACEPARENT_HEADER, "")
    if traceparent:
        parts = traceparent.split("-")
        if len(parts) >= 4:
            result["version"] = parts[0]
            # Normalize to lowercase per W3C spec, then validate hex format
            trace_id = parts[1].lower()
            span_id = parts[2].lower()
            if _is_valid_hex_id(trace_id, 32):
                result["trace_id"] = trace_id
            if _is_valid_hex_id(span_id, 16):
                result["parent_span_id"] = span_id
            result["flags"] = parts[3]

    # Parse tracestate - W3C format: "vendor1=value1,vendor2=value2"
    tracestate = headers.get(TRACESTATE_HEADER, "")
    if tracestate:
        # Split by ',' to get vendor entries
        for vendor_entry in tracestate.split(","):
            vendor_entry = vendor_entry.strip()
            if vendor_entry.startswith("risicare="):
                # Extract Risicare state (uses ';' as internal delimiter)
                risicare_state = vendor_entry[9:]  # len("risicare=")
                # Split by ';' to get key-value pairs
                for kv in risicare_state.split(_RISICARE_STATE_DELIMITER):
                    kv = kv.strip()
                    if "=" in kv:
                        k, v = kv.split("=", 1)
                        k_stripped = k.strip()
                        v_stripped = v.strip()
                        # Filter empty keys/values from malformed tracestate
                        if k_stripped and v_stripped:
                            result[k_stripped] = v_stripped

    return result


# =============================================================================
# Convenience Functions
# =============================================================================

def patch_all() -> None:
    """
    Patch ThreadPoolExecutor, ProcessPoolExecutor, and asyncio for context propagation.

    Call this once at SDK initialization.
    """
    patch_executors()
    patch_process_executors()
    patch_asyncio()


def unpatch_all() -> None:
    """Restore all original implementations."""
    unpatch_executors()
    unpatch_process_executors()
    unpatch_asyncio()


def _reset_for_testing() -> None:
    """
    Reset all global state for testing.

    This is intended for test fixtures only. Do not use in production.
    Clears all context variables and the span registry.
    """
    # Clear span registry
    _span_registry.clear()

    # Reset context variables to defaults
    # Note: We can't truly reset contextvars, but we can set them to None
    # which is their default. This only affects the current context.
    try:
        _current_session.set(None)
    except LookupError:
        pass
    try:
        _current_agent.set(None)
    except LookupError:
        pass
    try:
        _current_span.set(None)
    except LookupError:
        pass
    try:
        _current_phase.set(None)
    except LookupError:
        pass

    # Unpatch if patched
    unpatch_all()
