"""Enhanced OpenTelemetry tracing module for Robin AI.

This module provides a comprehensive, production-ready tracing system that consolidates
and enhances the existing tracing functionality with proper resource management,
circuit breaker protection, and correlation support.
"""

import asyncio
import os
import time
from collections.abc import AsyncGenerator, Generator, Sequence
from contextlib import asynccontextmanager, contextmanager
from functools import wraps
from typing import Any, Protocol

from opentelemetry import trace
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExportResult
from opentelemetry.sdk.trace.export import SpanExporter as OTLPSpanExporterProtocol
from opentelemetry.trace import SpanKind, Status, StatusCode

try:
    import grpc

    COMPRESSION_GZIP: Any = grpc.Compression.Gzip
except (ImportError, AttributeError):
    COMPRESSION_GZIP = None

# Updated imports to robin_commons
from robin_commons.log import logger
from robin_commons.resilience.breaker import CircuitBreaker
from robin_commons.telemetry.config import get_observability_config


class TracingError(Exception):
    """Base exception for tracing-related errors."""

    pass


class TracingConfigurationError(TracingError):
    """Raised when tracing configuration is invalid."""

    pass


class CircuitBreakerSpanExporter(OTLPSpanExporterProtocol):
    """OTLP span exporter with circuit breaker protection."""

    def __init__(
        self, otlp_exporter: OTLPSpanExporter, circuit_breaker: CircuitBreaker
    ) -> None:
        """Initialize span exporter with circuit breaker protection."""
        self._otlp_exporter = otlp_exporter
        self._circuit_breaker = circuit_breaker

    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
        """Export spans with circuit breaker protection."""
        if not self._circuit_breaker.can_execute():
            logger.warning(
                "Circuit breaker blocking span export",
                spans_count=len(spans),
                state=self._circuit_breaker.get_state().value,
            )
            return SpanExportResult.FAILURE

        try:
            result = self._otlp_exporter.export(spans)

            if result == SpanExportResult.SUCCESS:
                self._circuit_breaker.record_success()
                logger.debug(
                    "Span export succeeded",
                    spans_count=len(spans),
                    circuit_breaker_state=self._circuit_breaker.get_state().value,
                )
            else:
                self._circuit_breaker.record_failure()
                logger.warning(
                    "Span export failed",
                    spans_count=len(spans),
                    result=result,
                    circuit_breaker_failures=self._circuit_breaker.get_failure_count(),
                )

            return result

        except Exception as e:
            self._circuit_breaker.record_failure()
            logger.exception(
                "Span export failed with exception",
                spans_count=len(spans),
                error=str(e),
                circuit_breaker_failures=self._circuit_breaker.get_failure_count(),
            )
            return SpanExportResult.FAILURE

    def force_flush(self, timeout_millis: int = 30000) -> bool:
        """Force flush any buffered spans."""
        try:
            return self._otlp_exporter.force_flush(timeout_millis)
        except Exception as e:
            logger.warning(
                "Force flush failed", error=str(e), timeout_millis=timeout_millis
            )
            return False

    def shutdown(self) -> None:
        """Shut down the exporter."""
        try:
            self._otlp_exporter.shutdown()
        except Exception as e:
            logger.warning("Error during span exporter shutdown", error=str(e))


class SpanContextManager(Protocol):
    """Protocol for span context managers."""

    def __enter__(self) -> Any:
        """Enter span context."""
        ...

    def __exit__(
        self, exc_type: type | None, exc_val: Exception | None, exc_tb: Any
    ) -> None:
        """Exit span context."""
        ...


class AsyncSpanContextManager(Protocol):
    """Protocol for async span context managers."""

    async def __aenter__(self) -> Any:
        """Enter async span context."""
        ...

    async def __aexit__(
        self, exc_type: type | None, exc_val: Exception | None, exc_tb: Any
    ) -> None:
        """Exit async span context."""
        ...


class TracingService:
    """Enhanced tracing service with comprehensive observability features."""

    def __init__(self) -> None:
        """Initialize tracing service with configuration."""
        self.config = get_observability_config()
        self._tracer: trace.Tracer | None = None
        self._is_initialized = False

    def setup_tracing(self) -> None:
        """Initialize OpenTelemetry tracing with comprehensive configuration."""
        try:
            if not self.config.is_tracing_enabled():
                logger.info("OpenTelemetry tracing is disabled via feature flag")
                return

            # Create resource with comprehensive attributes
            resource_attrs = self.config.get_resource_attributes()
            resource = Resource.create(resource_attrs)
            provider = TracerProvider(resource=resource)

            # Configure exporter with circuit breaker protection
            exporter = self._create_span_exporter()
            batch_config = self.config.get_batch_export_config()

            processor = BatchSpanProcessor(
                exporter,
                max_queue_size=batch_config["max_queue_size"],
                max_export_batch_size=batch_config["max_export_batch_size"],
                export_timeout_millis=batch_config["export_timeout_millis"],
                schedule_delay_millis=batch_config["schedule_delay_millis"],
            )

            provider.add_span_processor(processor)
            trace.set_tracer_provider(provider)
            self._tracer = trace.get_tracer(self.config.get_service_name())

            self._is_initialized = True
            logger.info(
                "OpenTelemetry tracing initialized successfully",
                service_name=self.config.get_service_name(),
                alloy_enabled=self.config.is_alloy_enabled(),
            )

        except Exception as e:
            raise TracingConfigurationError(f"Failed to initialize tracing: {e}") from e

    def _create_span_exporter(self) -> CircuitBreakerSpanExporter:
        """Create span exporter with circuit breaker protection."""
        otlp_exporter = self._create_otlp_exporter()
        circuit_config = self.config.get_circuit_breaker_config()
        circuit_breaker = CircuitBreaker(circuit_config)

        logger.info(
            "Created span exporter with circuit breaker",
            failure_threshold=circuit_config.failure_threshold,
            recovery_timeout=circuit_config.recovery_timeout_seconds,
            success_threshold=circuit_config.success_threshold,
        )

        return CircuitBreakerSpanExporter(otlp_exporter, circuit_breaker)

    def _create_otlp_exporter(self) -> OTLPSpanExporter:
        """Create OTLP span exporter based on configuration."""
        endpoint = self.config.get_otlp_endpoint()
        headers = self.config.get_otlp_headers()
        protocol = self.config.get_otlp_protocol()

        logger.info(
            "Configuring OTLP exporter",
            endpoint=endpoint,
            protocol=protocol,
            alloy_enabled=self.config.is_alloy_enabled(),
            has_headers=bool(headers),
        )

        # Configure based on protocol and Alloy usage
        if self.config.is_alloy_enabled():
            # Local Alloy - no authentication needed, use compression if available
            if COMPRESSION_GZIP and protocol == "grpc":
                return OTLPSpanExporter(
                    endpoint=endpoint,
                    insecure=True,
                    compression=COMPRESSION_GZIP,
                )
            else:
                return OTLPSpanExporter(endpoint=endpoint, insecure=True)
        else:
            # Direct to Grafana Cloud - requires authentication
            if not headers:
                raise TracingConfigurationError(
                    "Headers required for direct Grafana Cloud export"
                )

            if COMPRESSION_GZIP and protocol == "grpc":
                return OTLPSpanExporter(
                    endpoint=endpoint,
                    headers=list(headers.items()),
                    compression=COMPRESSION_GZIP,
                )
            else:
                return OTLPSpanExporter(
                    endpoint=endpoint,
                    headers=list(headers.items()),
                )

    def create_span(
        self, name: str, kind: SpanKind = SpanKind.INTERNAL, **attributes: Any
    ) -> SpanContextManager:
        """Create a new span context manager."""
        if not self._is_initialized or not self._tracer:
            return _NoOpSpanContextManager()

        return _SpanContextManager(self._tracer, name, kind, attributes)

    def create_async_span(
        self, name: str, kind: SpanKind = SpanKind.INTERNAL, **attributes: Any
    ) -> AsyncSpanContextManager:
        """Create a new async span context manager."""
        if not self.config.is_advanced_features_enabled():
            return _NoOpAsyncSpanContextManager()

        if not self._is_initialized or not self._tracer:
            return _NoOpAsyncSpanContextManager()

        return _AsyncSpanContextManager(self._tracer, name, kind, attributes)


class _SpanContextManager:
    """Enhanced span context manager with correlation support."""

    def __init__(
        self,
        tracer: trace.Tracer,
        name: str,
        kind: SpanKind,
        attributes: dict[str, Any],
    ) -> None:
        self._tracer = tracer
        self._name = name
        self._kind = kind
        self._attributes = attributes
        self._span: trace.Span | None = None
        self._span_context_manager: Any = None
        self._start_time = 0.0

    def __enter__(self) -> trace.Span:
        """Enter span context and start timing."""
        # Create span and set it as current in one step
        span_attributes = {
            "span.type": self._kind.name.lower(),
            "service.operation": self._name,
            **self._attributes,
        }

        self._span_context_manager = self._tracer.start_as_current_span(
            self._name, kind=self._kind, attributes=span_attributes
        )

        # Enter the context manager which returns the span
        self._span = self._span_context_manager.__enter__()

        # Add correlation context to span and logs
        from robin_commons.telemetry.correlation import (
            add_span_attributes,
            set_trace_context,
        )

        add_span_attributes(self._span)
        set_trace_context()

        self._start_time = time.perf_counter()

        return self._span

    def __exit__(
        self, exc_type: type | None, exc_val: Exception | None, exc_tb: Any
    ) -> None:
        """Exit span context with proper status and duration."""
        if not self._span or not self._span_context_manager:
            return

        try:
            # Calculate and record duration
            duration_ms = (time.perf_counter() - self._start_time) * 1000
            self._span.set_attribute("duration_ms", duration_ms)

            # Set span status based on exception
            if exc_val:
                self._span.record_exception(exc_val)
                self._span.set_status(Status(StatusCode.ERROR, str(exc_val)))
                logger.exception(f"Error in {self._name}")
            else:
                self._span.set_status(Status(StatusCode.OK))

        finally:
            self._span_context_manager.__exit__(exc_type, exc_val, exc_tb)


class _AsyncSpanContextManager:
    """Enhanced async span context manager with correlation support."""

    def __init__(
        self,
        tracer: trace.Tracer,
        name: str,
        kind: SpanKind,
        attributes: dict[str, Any],
    ) -> None:
        self._tracer = tracer
        self._name = name
        self._kind = kind
        self._attributes = attributes
        self._span: trace.Span | None = None
        self._span_context_manager: Any = None
        self._start_time = 0.0

    async def __aenter__(self) -> trace.Span:
        """Enter async span context."""
        span_attributes = {
            "span.type": self._kind.name.lower(),
            "service.operation": self._name,
            "async": True,
            **self._attributes,
        }

        # Use start_as_current_span to both start and set current span
        self._span_context_manager = self._tracer.start_as_current_span(
            self._name, kind=self._kind, attributes=span_attributes
        )

        # Enter the context manager; it's fine to call synchronously here
        self._span = self._span_context_manager.__enter__()

        # Add correlation context
        from robin_commons.telemetry.correlation import (
            add_span_attributes,
            set_trace_context,
        )

        add_span_attributes(self._span)
        set_trace_context()

        self._start_time = time.perf_counter()

        return self._span

    async def __aexit__(
        self, exc_type: type | None, exc_val: Exception | None, exc_tb: Any
    ) -> None:
        """Exit async span context."""
        if not self._span or not self._span_context_manager:
            return

        try:
            # Calculate and record duration
            duration_ms = (time.perf_counter() - self._start_time) * 1000
            self._span.set_attribute("duration_ms", duration_ms)

            # Set span status
            if exc_val:
                self._span.record_exception(exc_val)
                self._span.set_status(Status(StatusCode.ERROR, str(exc_val)))
                logger.exception(f"Error in async {self._name}")
            else:
                self._span.set_status(Status(StatusCode.OK))

        finally:
            self._span_context_manager.__exit__(exc_type, exc_val, exc_tb)


class _NoOpSpanContextManager:
    """No-op span context manager for when tracing is disabled."""

    def __enter__(self) -> None:
        return None

    def __exit__(
        self, exc_type: type | None, exc_val: Exception | None, exc_tb: Any
    ) -> None:
        pass


class _NoOpAsyncSpanContextManager:
    """No-op async span context manager for when async spans are disabled."""

    async def __aenter__(self) -> None:
        return None

    async def __aexit__(
        self, exc_type: type | None, exc_val: Exception | None, exc_tb: Any
    ) -> None:
        pass


# Global tracing service instance
_tracing_service: TracingService | None = None


def setup_observability() -> None:
    """Initialize the complete observability system."""
    global _tracing_service

    # Skip during tests to avoid connection errors
    if os.getenv("PYTEST_CURRENT_TEST") or os.getenv("TESTING") == "true":
        logger.info("Observability disabled during tests")
        return

    _tracing_service = TracingService()
    _tracing_service.setup_tracing()


@contextmanager
def span(
    name: str, kind: SpanKind = SpanKind.INTERNAL, **attributes: Any
) -> Generator[trace.Span | None]:
    """Create a new span for tracing operations.

    Args:
        name: Descriptive name for the operation being traced
        kind: Type of span (INTERNAL for business logic, CLIENT for external calls, etc.)
        **attributes: Additional attributes to attach to the span

    Yields:
        OpenTelemetry Span object or None if tracing is disabled

    Example:
        with span("user_authentication", user_id=123, method="google_sso"):
            user = authenticate_user(token)
            logger.info("User authenticated", user_id=user.id)
    """
    if not _tracing_service:
        yield None
        return

    with _tracing_service.create_span(name, kind, **attributes) as span_obj:
        yield span_obj


@asynccontextmanager
async def async_span(
    name: str, kind: SpanKind = SpanKind.INTERNAL, **attributes: Any
) -> AsyncGenerator[trace.Span | None]:
    """Create a new async span for tracing async operations.

    Args:
        name: Descriptive name for the operation being traced
        kind: Type of span (INTERNAL for business logic, CLIENT for external calls, etc.)
        **attributes: Additional attributes to attach to the span

    Yields:
        OpenTelemetry Span object or None if tracing is disabled

    Example:
        async with async_span("database_query", table="users", query_type="select"):
            users = await db.get_users()
    """
    if not _tracing_service:
        yield None
        return

    async with _tracing_service.create_async_span(name, kind, **attributes) as span_obj:
        yield span_obj


def traced(
    name: str | None = None, kind: SpanKind = SpanKind.INTERNAL, **attributes: Any
) -> Any:
    """Decorator for automatic tracing of functions.

    Args:
        name: Optional custom name for the span. If None, uses module.function_name
        kind: Type of span
        **attributes: Additional attributes to attach to the span

    Returns:
        Decorated function with automatic tracing

    Example:
        @traced("database.get_user")
        async def get_user_by_id(user_id: int) -> User:
            return await db.get_user(user_id)
    """

    def decorator(func: Any) -> Any:
        operation_name = name or f"{func.__module__}.{func.__name__}"
        if asyncio.iscoroutinefunction(func):

            @wraps(func)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                # Check if advanced observability features are enabled at runtime
                config = get_observability_config()
                if not config.is_advanced_features_enabled():
                    return await func(*args, **kwargs)

                async with async_span(operation_name, kind, **attributes):
                    return await func(*args, **kwargs)

            return async_wrapper
        else:

            @wraps(func)
            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                # Check if advanced observability features are enabled at runtime
                config = get_observability_config()
                if not config.is_advanced_features_enabled():
                    return func(*args, **kwargs)

                with span(operation_name, kind, **attributes):
                    return func(*args, **kwargs)

            return sync_wrapper

    return decorator


def add_span_event(
    name: str,
    attributes: dict[str, Any] | None = None,
) -> None:
    """Add an event to the current active span (if any).

    This is useful for adding lightweight events without creating separate spans,
    keeping traces clean and contextual.

    Args:
        name: Event name (e.g., "nats.publish.cmd.task.create")
        attributes: Optional attributes to attach to the event

    Example:
        # Within an existing traced operation
        add_span_event(
            "nats.publish.cmd.task.create",
            attributes={"stream": "USER_COMMANDS", "task_id": "task-123"}
        )
    """
    current_span = trace.get_current_span()
    if current_span and current_span.is_recording():
        current_span.add_event(name, attributes=attributes or {})


# Export public interface
__all__ = [
    "setup_observability",
    "span",
    "async_span",
    "traced",
    "add_span_event",
    "TracingError",
    "TracingConfigurationError",
    "SpanKind",
]
