"""Enhanced HTTP tracing middleware for Robin AI.

This middleware provides comprehensive HTTP request tracing with OpenTelemetry,
including proper W3C trace context propagation, correlation ID management,
and performance monitoring.
"""

from collections.abc import Awaitable, Callable
from typing import Any, cast

from fastapi import Request, Response
from opentelemetry import trace
from opentelemetry.propagate import extract, inject
from opentelemetry.trace import SpanKind
from starlette.middleware.base import BaseHTTPMiddleware

from robin_commons.log import logger
from robin_commons.telemetry.correlation import (
    ensure_correlation_id,
    set_correlation_from_headers,
    set_request_id,
    set_trace_context,
    set_user_id,
)


class TraceMiddleware(BaseHTTPMiddleware):
    """Enhanced HTTP request tracing middleware with comprehensive observability.

    Features:
    - W3C Trace Context propagation
    - Correlation ID management
    - Request/response performance monitoring
    - User context extraction
    - Proper error handling and status codes
    """

    async def dispatch(
        self, request: Request, call_next: Callable[[Request], Awaitable[Response]]
    ) -> Response:
        """Process HTTP request with enhanced distributed tracing.

        Args:
            request: Incoming HTTP request
            call_next: Next middleware/handler in the chain

        Returns:
            HTTP response with trace headers and correlation context
        """
        # Extract trace context from incoming headers (W3C Trace Context)
        context = extract(request.headers)

        # Set up correlation context from headers
        self._setup_correlation_context(request)

        # Create span for this HTTP request
        tracer = trace.get_tracer(__name__)
        operation_name = f"{request.method} {self._get_route_pattern(request)}"

        with tracer.start_as_current_span(
            operation_name,
            context=context,
            kind=SpanKind.SERVER,
            attributes=self._get_request_attributes(request),
        ) as span:
            # Store trace context in request state and correlation
            self._setup_trace_context(request, span)

            try:
                # Process request
                response = await call_next(request)

                # Add HTTP response attributes
                self._add_response_attributes(span, response)

                # Set span status based on HTTP status code
                self._set_span_status(span, response.status_code)

                # Add correlation and trace headers to response
                self._add_response_headers(request, response)

                return response

            except Exception as e:
                # Record exception in span
                span.record_exception(e)
                span.set_status(trace.Status(trace.StatusCode.ERROR, str(e)))

                # Log error with full context
                logger.error(
                    f"HTTP request failed: {operation_name}",
                    error=str(e),
                    method=request.method,
                    url=str(request.url),
                    trace_id=getattr(request.state, "trace_id", None),
                    span_id=getattr(request.state, "span_id", None),
                    correlation_id=getattr(request.state, "correlation_id", None),
                    exc_info=True,
                )

                raise

    def _setup_correlation_context(self, request: Request) -> None:
        """Set up correlation context from request headers."""
        # Extract correlation context from headers
        set_correlation_from_headers(dict(request.headers))

        # Ensure we have a correlation ID
        correlation_id = ensure_correlation_id()
        request.state.correlation_id = correlation_id

        # Generate request ID if not provided
        request_id = request.headers.get("x-request-id")
        if not request_id:
            import uuid

            request_id = str(uuid.uuid4())
        set_request_id(request_id)
        request.state.request_id = request_id

        # Extract user context from headers (set by APISix or auth middleware)
        if user_id := request.headers.get("x-user-id"):
            set_user_id(user_id)
            request.state.user_id = user_id

    def _setup_trace_context(self, request: Request, span: trace.Span) -> None:
        """Set up trace context in request state and correlation."""
        span_context = span.get_span_context()
        trace_id = format(span_context.trace_id, "032x")
        span_id = format(span_context.span_id, "016x")

        # Store in request state for application use
        request.state.trace_id = trace_id
        request.state.span_id = span_id

        # Set in correlation context for logging
        set_trace_context(trace_id, span_id)

    def _get_route_pattern(self, request: Request) -> str:
        """Get the route pattern for the request."""
        # Try to get the route pattern from FastAPI
        if hasattr(request, "route") and hasattr(request.route, "path"):
            return cast(str, request.route.path)

        # Fallback to URL path
        return request.url.path

    def _get_request_attributes(self, request: Request) -> dict[str, Any]:
        """Get comprehensive request attributes for the span."""
        attributes = {
            # Standard HTTP attributes
            "http.method": request.method,
            "http.url": str(request.url),
            "http.scheme": request.url.scheme,
            "http.host": request.url.hostname or "unknown",
            "http.target": request.url.path,
            "http.user_agent": request.headers.get("user-agent", "unknown"),
            "http.route": self._get_route_pattern(request),
            # Request metadata
            "http.request_content_length": request.headers.get("content-length", 0),
            "http.request_content_type": request.headers.get("content-type", "unknown"),
            # Client information
            "client.address": self._get_client_ip(request),
            # Robin AI specific attributes
            "service.component": "http-server",
            "service.operation": f"{request.method} {self._get_route_pattern(request)}",
        }

        # Add query parameters (sanitized)
        if request.query_params:
            # Only include non-sensitive query parameters
            safe_params = {}
            for key, value in request.query_params.items():
                if not self._is_sensitive_param(key):
                    safe_params[key] = value
            if safe_params:
                attributes["http.query_params"] = str(safe_params)

        # Add correlation context
        if correlation_id := getattr(request.state, "correlation_id", None):
            attributes["correlation.id"] = correlation_id

        if request_id := getattr(request.state, "request_id", None):
            attributes["correlation.request_id"] = request_id

        if user_id := getattr(request.state, "user_id", None):
            attributes["correlation.user_id"] = user_id

        return attributes

    def _add_response_attributes(self, span: trace.Span, response: Response) -> None:
        """Add response attributes to the span."""
        attributes: dict[str, str | int] = {
            "http.status_code": response.status_code,
            "http.response_content_type": response.headers.get(
                "content-type", "unknown"
            ),
        }

        # Add response size if available
        if content_length := response.headers.get("content-length"):
            attributes["http.response_content_length"] = int(content_length)
        elif hasattr(response, "body") and response.body:
            attributes["http.response_content_length"] = len(response.body)

        span.set_attributes(attributes)

    def _set_span_status(self, span: trace.Span, status_code: int) -> None:
        """Set span status based on HTTP status code."""
        if status_code >= 400:
            if status_code >= 500:
                span.set_status(
                    trace.Status(trace.StatusCode.ERROR, f"HTTP {status_code}")
                )
            else:
                # 4xx errors are client errors, not server errors
                span.set_status(trace.Status(trace.StatusCode.OK))
        else:
            span.set_status(trace.Status(trace.StatusCode.OK))

    def _add_response_headers(self, request: Request, response: Response) -> None:
        """Add correlation and trace headers to the response."""
        # Add trace context headers
        if trace_id := getattr(request.state, "trace_id", None):
            response.headers["X-Trace-Id"] = trace_id

        if span_id := getattr(request.state, "span_id", None):
            response.headers["X-Span-Id"] = span_id

        # Add correlation headers
        if correlation_id := getattr(request.state, "correlation_id", None):
            response.headers["X-Correlation-ID"] = correlation_id

        if request_id := getattr(request.state, "request_id", None):
            response.headers["X-Request-ID"] = request_id

        # Inject W3C trace context for downstream services
        inject(response.headers)

    def _get_client_ip(self, request: Request) -> str:
        """Get the client IP address from the request."""
        # Check for forwarded headers (from load balancers/proxies)
        forwarded_for = request.headers.get("x-forwarded-for")
        if forwarded_for:
            # Take the first IP in the chain
            return forwarded_for.split(",")[0].strip()

        real_ip = request.headers.get("x-real-ip")
        if real_ip:
            return real_ip

        # Fallback to client host
        if hasattr(request, "client") and request.client:
            return request.client.host

        return "unknown"

    def _is_sensitive_param(self, param_name: str) -> bool:
        """Check if a query parameter contains sensitive information."""
        sensitive_params = {
            "password",
            "token",
            "key",
            "secret",
            "auth",
            "api_key",
            "access_token",
            "refresh_token",
            "session",
            "cookie",
        }
        param_lower = param_name.lower()
        return any(sensitive in param_lower for sensitive in sensitive_params)
