# Copyright (c) Metis. All rights reserved.

"""Unified tracing context for MantisDK.

This module provides context variables and context managers for ALL tracing
metadata: call_type, tags, environment. The TracingContextSpanProcessor
reads from these and injects into spans.

Run/rollout context (run_id, session_id) is read directly from the run module.

Usage::

    import mantisdk.tracing as tracing

    # Context manager for ad-hoc tagging
    with tracing.tracing_context(call_type="agent-call", tags=["gepa", "train"]):
        response = client.chat.completions.create(...)  # Tagged automatically

    # Decorator for consistent function tagging
    @tracing.call_type_decorator("judge-call")
    def my_judge_function():
        response = client.chat.completions.create(...)  # Tagged as judge-call
"""

from __future__ import annotations

import asyncio
from contextlib import contextmanager
from contextvars import ContextVar
from functools import wraps
from typing import Any, Callable, Generator, List, Optional, TypeVar

# Type variable for decorators
F = TypeVar("F", bound=Callable[..., Any])

# =============================================================================
# Context Variables for Tracing Metadata
# =============================================================================

_call_type_context: ContextVar[Optional[str]] = ContextVar("mantis_call_type", default=None)
_tags_context: ContextVar[Optional[List[str]]] = ContextVar("mantis_tags", default=None)
_environment_context: ContextVar[Optional[str]] = ContextVar("mantis_environment", default=None)
# Proxy-specific session context. Used by RolloutAttemptMiddleware in llm_proxy
# to set the session_id for proxy spans (where there's no Run/Rollout object).
_proxy_session_context: ContextVar[Optional[str]] = ContextVar("mantis_proxy_session", default=None)


# =============================================================================
# Getters
# =============================================================================


def get_current_call_type() -> Optional[str]:
    """Get the current LLM call type (e.g., 'agent-call', 'judge-call').

    Returns:
        The current call type string, or None if not set.
    """
    return _call_type_context.get()


def get_current_tags() -> Optional[List[str]]:
    """Get the current tracing tags.

    Returns:
        The current list of tags, or None if not set.
    """
    return _tags_context.get()


def get_current_environment() -> Optional[str]:
    """Get the current tracing environment.

    Returns:
        The current environment string, or None if not set.
    """
    return _environment_context.get()


def get_proxy_session_id() -> Optional[str]:
    """Get the proxy-specific session ID.

    This is set by RolloutAttemptMiddleware in the LLM proxy when there's
    no Run/Rollout object in scope. Falls back to None, in which case
    TracingContextSpanProcessor won't set session_id on proxy spans.

    Returns:
        The proxy session ID string (typically the rollout_id), or None.
    """
    return _proxy_session_context.get()


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


@contextmanager
def tracing_context(
    *,
    call_type: Optional[str] = None,
    tags: Optional[List[str]] = None,
    environment: Optional[str] = None,
    session_id: Optional[str] = None,
) -> Generator[None, None, None]:
    """Set tracing context for the duration of the block.

    All spans created within this block will be tagged with the provided
    metadata via the TracingContextSpanProcessor.

    Args:
        call_type: LLM call type (e.g., "agent-call", "judge-call", "reflection-call")
        tags: List of tags for the traces
        environment: Environment name (e.g., "mantisdk-gepa")
        session_id: Session ID for proxy spans (set by RolloutAttemptMiddleware)

    Example::

        with tracing_context(call_type="agent-call", tags=["gepa", "train"]):
            response = client.chat.completions.create(...)  # Tagged automatically
    """
    tokens: List[tuple[str, Any]] = []

    if call_type is not None:
        tokens.append(("call_type", _call_type_context.set(call_type)))
    if tags is not None:
        tokens.append(("tags", _tags_context.set(tags)))
    if environment is not None:
        tokens.append(("environment", _environment_context.set(environment)))
    if session_id is not None:
        tokens.append(("session_id", _proxy_session_context.set(session_id)))

    try:
        yield
    finally:
        for name, token in tokens:
            if name == "call_type":
                _call_type_context.reset(token)
            elif name == "tags":
                _tags_context.reset(token)
            elif name == "environment":
                _environment_context.reset(token)
            elif name == "session_id":
                _proxy_session_context.reset(token)


# =============================================================================
# Decorators
# =============================================================================


def call_type_decorator(call_type: str) -> Callable[[F], F]:
    """Factory to create call-type tagging decorators.

    Creates a decorator that sets the call_type context for the duration
    of the decorated function. All spans created within the function
    will have the `mantis.call_type` attribute set.

    Args:
        call_type: The call type to set (e.g., "agent-call", "judge-call")

    Returns:
        A decorator that wraps functions with tracing context.

    Example::

        # Create reusable decorators
        agent = call_type_decorator("agent-call")
        judge = call_type_decorator("judge-call")

        @agent
        def my_agent_function():
            # LLM calls here are tagged as "agent-call"
            response = client.chat.completions.create(...)

        @judge
        def my_judge_function():
            # LLM calls here are tagged as "judge-call"
            response = client.chat.completions.create(...)
    """

    def decorator(fn: F) -> F:
        if asyncio.iscoroutinefunction(fn):

            @wraps(fn)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                with tracing_context(call_type=call_type):
                    return await fn(*args, **kwargs)

            return async_wrapper  # type: ignore[return-value]
        else:

            @wraps(fn)
            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                with tracing_context(call_type=call_type):
                    return fn(*args, **kwargs)

            return sync_wrapper  # type: ignore[return-value]

    return decorator
