"""OpenTelemetry tracing configuration.

Implements distributed tracing for:
- API requests
- Database queries
- External API calls
- Background tasks
"""

import asyncio
from typing import Optional, Callable
from functools import wraps
from contextlib import contextmanager

from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.sdk.resources import Resource, SERVICE_NAME, SERVICE_VERSION
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.exporter.jaeger.thrift import JaegerExporter
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
from opentelemetry.instrumentation.redis import RedisInstrumentor
from opentelemetry.instrumentation.celery import CeleryInstrumentor
from opentelemetry.trace import Status, StatusCode

from src.core.config import settings


# Global tracer provider
_tracer_provider: Optional[TracerProvider] = None
_tracer: Optional[trace.Tracer] = None


def setup_tracing(
    service_name: str = "mockupAWS",
    service_version: str = "1.0.0",
    jaeger_endpoint: Optional[str] = None,
    otlp_endpoint: Optional[str] = None,
) -> TracerProvider:
    """Setup OpenTelemetry tracing.

    Args:
        service_name: Name of the service
        service_version: Version of the service
        jaeger_endpoint: Jaeger collector endpoint
        otlp_endpoint: OTLP collector endpoint

    Returns:
        Configured TracerProvider
    """
    global _tracer_provider, _tracer

    # Create resource
    resource = Resource.create(
        {
            SERVICE_NAME: service_name,
            SERVICE_VERSION: service_version,
            "deployment.environment": "production"
            if not getattr(settings, "debug", False)
            else "development",
        }
    )

    # Create tracer provider
    _tracer_provider = TracerProvider(resource=resource)

    # Add exporters
    if jaeger_endpoint or getattr(settings, "jaeger_endpoint", None):
        jaeger_exporter = JaegerExporter(
            agent_host_name=jaeger_endpoint
            or getattr(settings, "jaeger_endpoint", "localhost"),
            agent_port=getattr(settings, "jaeger_port", 6831),
        )
        _tracer_provider.add_span_processor(BatchSpanProcessor(jaeger_exporter))

    if otlp_endpoint or getattr(settings, "otlp_endpoint", None):
        otlp_exporter = OTLPSpanExporter(
            endpoint=otlp_endpoint or getattr(settings, "otlp_endpoint"),
        )
        _tracer_provider.add_span_processor(BatchSpanProcessor(otlp_exporter))

    # Set as global provider
    trace.set_tracer_provider(_tracer_provider)

    # Get tracer
    _tracer = trace.get_tracer(service_name, service_version)

    return _tracer_provider


def instrument_fastapi(app) -> None:
    """Instrument FastAPI application for tracing.

    Args:
        app: FastAPI application instance
    """
    FastAPIInstrumentor.instrument_app(
        app,
        tracer_provider=_tracer_provider,
    )


def instrument_sqlalchemy(engine) -> None:
    """Instrument SQLAlchemy for database query tracing.

    Args:
        engine: SQLAlchemy engine instance
    """
    SQLAlchemyInstrumentor().instrument(
        engine=engine,
        tracer_provider=_tracer_provider,
    )


def instrument_redis() -> None:
    """Instrument Redis for caching operation tracing."""
    RedisInstrumentor().instrument(tracer_provider=_tracer_provider)


def instrument_celery() -> None:
    """Instrument Celery for task tracing."""
    CeleryInstrumentor().instrument(tracer_provider=_tracer_provider)


def get_tracer() -> trace.Tracer:
    """Get the global tracer.

    Returns:
        Tracer instance
    """
    if _tracer is None:
        raise RuntimeError("Tracing not initialized. Call setup_tracing() first.")
    return _tracer


@contextmanager
def start_span(
    name: str,
    kind: trace.SpanKind = trace.SpanKind.INTERNAL,
    attributes: Optional[dict] = None,
):
    """Context manager for starting a span.

    Args:
        name: Span name
        kind: Span kind
        attributes: Span attributes

    Yields:
        Span context
    """
    tracer = get_tracer()
    with tracer.start_as_current_span(name, kind=kind) as span:
        if attributes:
            for key, value in attributes.items():
                span.set_attribute(key, value)
        yield span


def trace_function(
    name: Optional[str] = None,
    attributes: Optional[dict] = None,
):
    """Decorator to trace function execution.

    Args:
        name: Span name (defaults to function name)
        attributes: Additional span attributes

    Returns:
        Decorated function
    """

    def decorator(func: Callable) -> Callable:
        span_name = name or func.__name__

        @wraps(func)
        async def async_wrapper(*args, **kwargs):
            tracer = get_tracer()
            with tracer.start_as_current_span(span_name) as span:
                # Add function attributes
                span.set_attribute("function.name", func.__name__)
                span.set_attribute("function.module", func.__module__)

                if attributes:
                    for key, value in attributes.items():
                        span.set_attribute(key, value)

                try:
                    result = await func(*args, **kwargs)
                    span.set_status(Status(StatusCode.OK))
                    return result
                except Exception as e:
                    span.set_status(Status(StatusCode.ERROR, str(e)))
                    span.record_exception(e)
                    raise

        @wraps(func)
        def sync_wrapper(*args, **kwargs):
            tracer = get_tracer()
            with tracer.start_as_current_span(span_name) as span:
                span.set_attribute("function.name", func.__name__)
                span.set_attribute("function.module", func.__module__)

                if attributes:
                    for key, value in attributes.items():
                        span.set_attribute(key, value)

                try:
                    result = func(*args, **kwargs)
                    span.set_status(Status(StatusCode.OK))
                    return result
                except Exception as e:
                    span.set_status(Status(StatusCode.ERROR, str(e)))
                    span.record_exception(e)
                    raise

        return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper

    return decorator


def trace_db_query(operation: str, table: str):
    """Decorator to trace database queries.

    Args:
        operation: Query operation (SELECT, INSERT, etc.)
        table: Table name

    Returns:
        Decorator function
    """
    return trace_function(
        name=f"db.query.{table}.{operation}",
        attributes={
            "db.operation": operation,
            "db.table": table,
        },
    )


def trace_external_call(service: str, operation: str):
    """Decorator to trace external API calls.

    Args:
        service: External service name
        operation: Operation being performed

    Returns:
        Decorator function
    """
    return trace_function(
        name=f"external.{service}.{operation}",
        attributes={
            "external.service": service,
            "external.operation": operation,
        },
    )


class TracingMiddleware:
    """FastAPI middleware for request tracing with correlation."""

    def __init__(self, app):
        self.app = app

    async def __call__(self, scope, receive, send):
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        from fastapi import Request

        request = Request(scope, receive)
        tracer = get_tracer()

        # Extract or create trace context
        with tracer.start_as_current_span(
            f"{request.method} {request.url.path}",
            kind=trace.SpanKind.SERVER,
        ) as span:
            # Add request attributes
            span.set_attribute("http.method", request.method)
            span.set_attribute("http.url", str(request.url))
            span.set_attribute("http.route", request.url.path)
            span.set_attribute("http.host", request.headers.get("host", "unknown"))
            span.set_attribute(
                "http.user_agent", request.headers.get("user-agent", "unknown")
            )

            # Add correlation ID if present
            correlation_id = request.headers.get("x-correlation-id")
            if correlation_id:
                span.set_attribute("correlation.id", correlation_id)

            try:
                await self.app(scope, receive, send)
                span.set_status(Status(StatusCode.OK))
            except Exception as e:
                span.set_status(Status(StatusCode.ERROR, str(e)))
                span.record_exception(e)
                raise
