# Copyright (c) Microsoft. All rights reserved.

"""Tracing configuration for algorithms.

This module provides TracingConfig -- a standardized way for algorithms to
define their tracing metadata (environment and tags) that flows through
to Mantis/Insight.

For call-type tagging and context management, use the unified tracing
context system::

    from mantisdk.tracing.context import call_type_decorator, get_current_call_type
    # or
    import mantisdk.tracing as tracing
    tracing.call_type_decorator("my-call-type")
"""

import uuid
from dataclasses import dataclass, field
from typing import List, Optional

# Re-export from canonical location for backward compatibility
from mantisdk.tracing.context import (
    call_type_decorator,
    get_current_call_type,
)


def inject_call_type_header(extra_headers: Optional[dict] = None) -> dict:
    """Inject x-mantis-call-type header from context into extra_headers.

    This reads the current call type from context (set by decorators like
    @gepa.agent, @gepa.judge) and adds it to extra_headers for LLM calls.
    The LLM proxy receives this header and tags the trace accordingly.

    Args:
        extra_headers: Optional existing headers dict to merge into.

    Returns:
        A dict with x-mantis-call-type header if call type is set.

    Example:
        >>> @gepa.agent
        >>> def my_agent(client):
        ...     return client.chat.completions.create(
        ...         model="gpt-4",
        ...         messages=[...],
        ...         extra_headers=inject_call_type_header()
        ...     )
    """
    if extra_headers is None:
        extra_headers = {}

    call_type = get_current_call_type()
    if call_type:
        extra_headers["x-mantis-call-type"] = call_type

    return extra_headers


@dataclass
class TracingConfig:
    """Configuration for algorithm-specific tracing metadata.

    Each algorithm (GEPA, VERL, etc.) should define its own TracingConfig
    to identify its traces in Mantis. The environment, tags, and session_id are
    propagated through the LLM proxy and stored with traces.

    Example:
        >>> config = TracingConfig(
        ...     environment="mantisdk-gepa",
        ...     algorithm_name="gepa"
        ... )
        >>> config.get_tags("train")
        ['gepa', 'train']

    Attributes:
        environment: The environment name for traces (e.g., "mantisdk-gepa").
            This appears in Mantis's environment column.
        algorithm_name: The algorithm identifier used as the base tag.
        session_id: Unique session identifier for grouping related traces.
            Auto-generated if not provided.
    """

    environment: str
    algorithm_name: str
    session_id: str = field(default_factory=lambda: f"session-{uuid.uuid4().hex[:12]}")

    def get_tags(self, trace_type: str) -> List[str]:
        """Generate tags for a specific trace type.

        Args:
            trace_type: The type of trace (e.g., "train", "validation",
                "reflection", "test").

        Returns:
            A list of tags: [algorithm_name, trace_type]
        """
        return [self.algorithm_name, trace_type]

    def get_detailed_tags(
        self,
        phase: str,
        extra_tags: Optional[List[str]] = None,
    ) -> List[str]:
        """Generate tags with execution context.

        Args:
            phase: Execution phase (e.g., "train", "validation", "reflection").
            extra_tags: Optional algorithm-specific tags to append.

        Returns:
            A list of tags, e.g.: ["gepa", "train", "custom-tag"]
        """
        tags = [self.algorithm_name, phase]

        # Add any algorithm-specific extra tags
        if extra_tags:
            tags.extend(extra_tags)

        return tags

    def to_metadata(self, trace_type: Optional[str] = None) -> dict:
        """Convert to a metadata dict for rollouts or LLM calls.

        Args:
            trace_type: Optional trace type to include in tags.
                If None, only the algorithm_name is included in tags.

        Returns:
            Dict with "environment", "tags", and "session_id" keys.
        """
        if trace_type:
            tags = self.get_tags(trace_type)
        else:
            tags = [self.algorithm_name]

        return {
            "environment": self.environment,
            "tags": tags,
            "session_id": self.session_id,
        }

    def to_detailed_metadata(
        self,
        phase: str,
        extra_tags: Optional[List[str]] = None,
    ) -> dict:
        """Convert to a metadata dict with execution context.

        Args:
            phase: Execution phase (e.g., "train", "validation", "reflection").
            extra_tags: Optional algorithm-specific tags to append.

        Returns:
            Dict with "environment", "tags", and "session_id" keys.
        """
        return {
            "environment": self.environment,
            "tags": self.get_detailed_tags(phase, extra_tags),
            "session_id": self.session_id,
        }
