"""Distributed tracing using AWS Lambda Powertools and X-Ray.

This module provides X-Ray tracing integration for understanding
request flows through the AI agent system.

Usage:
    from agent_factory.observability import get_tracer

    tracer = get_tracer()

    @tracer.capture_method
    def my_function():
        # Will be traced as a subsegment
        ...

    with tracer.provider.in_subsegment("custom_segment") as subsegment:
        subsegment.put_annotation("key", "value")
        ...

For non-Lambda environments, provides a no-op fallback.
"""

import os
from functools import lru_cache
from typing import Any, Callable, Optional, TypeVar

# Try to import Powertools Tracer
try:
    from aws_lambda_powertools import Tracer as PowertoolsTracer

    HAS_POWERTOOLS = True
except ImportError:
    HAS_POWERTOOLS = False
    PowertoolsTracer = None


F = TypeVar("F", bound=Callable[..., Any])


class NoOpSubsegment:
    """No-op subsegment for non-Lambda environments."""

    def put_annotation(self, key: str, value: Any) -> None:
        """No-op annotation."""
        pass

    def put_metadata(self, key: str, value: Any, namespace: str = "default") -> None:
        """No-op metadata."""
        pass

    def __enter__(self) -> "NoOpSubsegment":
        return self

    def __exit__(self, *args: Any) -> None:
        pass


class NoOpProvider:
    """No-op provider for non-Lambda environments."""

    def in_subsegment(self, name: str) -> NoOpSubsegment:
        """Return no-op subsegment."""
        return NoOpSubsegment()

    def in_subsegment_async(self, name: str) -> NoOpSubsegment:
        """Return no-op async subsegment."""
        return NoOpSubsegment()

    def put_annotation(self, key: str, value: Any) -> None:
        """No-op annotation."""
        pass

    def put_metadata(self, key: str, value: Any, namespace: str = "default") -> None:
        """No-op metadata."""
        pass


class NoOpTracer:
    """No-op tracer class for non-Lambda environments."""

    def __init__(self, service: Optional[str] = None):
        self.service = service
        self.provider = NoOpProvider()

    def capture_method(self, method: F) -> F:
        """No-op method capture - just returns the method unchanged."""
        return method

    def capture_lambda_handler(self, handler: F) -> F:
        """No-op lambda handler capture."""
        return handler

    def put_annotation(self, key: str, value: Any) -> None:
        """No-op annotation."""
        pass

    def put_metadata(self, key: str, value: Any, namespace: str = "default") -> None:
        """No-op metadata."""
        pass


@lru_cache(maxsize=8)
def get_tracer(service: Optional[str] = None) -> Any:
    """Get a configured tracer instance for X-Ray tracing.

    Uses AWS Lambda Powertools Tracer when available, which integrates
    with AWS X-Ray for distributed tracing. Falls back to a no-op
    implementation in non-Lambda environments.

    Args:
        service: Service name (defaults to env or "agent")

    Returns:
        Tracer instance (Powertools or NoOp)

    Example:
        tracer = get_tracer()

        @tracer.capture_method
        def process_message():
            tracer.put_annotation("message_type", "user")
            ...
    """
    service_name = service or os.getenv("POWERTOOLS_SERVICE_NAME", "agent")

    if HAS_POWERTOOLS:
        return PowertoolsTracer(service=service_name)
    else:
        return NoOpTracer(service=service_name)


def trace_tool_execution(tracer: Any, tool_name: str) -> Any:
    """Context manager for tracing tool execution.

    Args:
        tracer: Tracer instance
        tool_name: Name of the tool being executed

    Returns:
        Context manager for the trace segment

    Example:
        with trace_tool_execution(tracer, "search_docs") as segment:
            result = tool.execute(**params)
            segment.put_metadata("result", result)
    """
    if HAS_POWERTOOLS:
        return tracer.provider.in_subsegment(f"Tool: {tool_name}")
    else:
        return tracer.provider.in_subsegment(tool_name)


def trace_model_call(tracer: Any, model_id: str) -> Any:
    """Context manager for tracing AI model calls.

    Args:
        tracer: Tracer instance
        model_id: ID of the model being called

    Returns:
        Context manager for the trace segment

    Example:
        with trace_model_call(tracer, "claude-3-5-sonnet") as segment:
            segment.put_annotation("message_count", len(messages))
            response = model.converse(messages)
    """
    if HAS_POWERTOOLS:
        return tracer.provider.in_subsegment(f"Model: {model_id}")
    else:
        return tracer.provider.in_subsegment(model_id)
