"""Distributed tracing support for HTTP requests.

Provides OpenTelemetry-compatible tracing with:
- Trace ID propagation via headers
- Span creation for HTTP requests
- Context propagation to downstream services
- Performance metrics
"""

from __future__ import annotations

import contextvars
import logging
import time
import uuid
from typing import Any

from aiohttp import web

logger = logging.getLogger(__name__)

# Context variables for trace propagation
_trace_id_var: contextvars.ContextVar[str | None] = contextvars.ContextVar("albus_trace_id", default=None)
_span_id_var: contextvars.ContextVar[str | None] = contextvars.ContextVar("albus_span_id", default=None)
_parent_span_id_var: contextvars.ContextVar[str | None] = contextvars.ContextVar(
    "albus_parent_span_id", default=None
)


def get_trace_id() -> str | None:
    """Get current trace ID."""
    return _trace_id_var.get()


def get_span_id() -> str | None:
    """Get current span ID."""
    return _span_id_var.get()


def set_trace_context(trace_id: str, span_id: str, parent_span_id: str | None = None) -> None:
    """Set trace context."""
    _trace_id_var.set(trace_id)
    _span_id_var.set(span_id)
    if parent_span_id:
        _parent_span_id_var.set(parent_span_id)


def generate_trace_id() -> str:
    """Generate a new trace ID (W3C TraceContext format)."""
    # W3C TraceContext: 32 hex characters (16 bytes)
    return uuid.uuid4().hex


def generate_span_id() -> str:
    """Generate a new span ID (W3C TraceContext format)."""
    # W3C TraceContext: 16 hex characters (8 bytes)
    return uuid.uuid4().hex[:16]


@web.middleware
async def tracing_middleware(request: web.Request, handler: Any) -> web.Response:
    """Distributed tracing middleware for HTTP requests.

    Supports W3C TraceContext format:
    - traceparent: 00-{trace_id}-{span_id}-{flags}
    - tracestate: Optional vendor-specific data

    Creates a span for each HTTP request and propagates trace context.
    """
    # Extract trace context from headers (W3C TraceContext)
    traceparent = request.headers.get("traceparent")
    trace_id = None
    parent_span_id = None

    if traceparent:
        # Parse W3C TraceContext format: 00-{trace_id}-{span_id}-{flags}
        parts = traceparent.split("-")
        if len(parts) >= 4 and parts[0] == "00":
            trace_id = parts[1]
            parent_span_id = parts[2]

    # Generate new trace ID if not present
    if not trace_id:
        trace_id = generate_trace_id()

    # Generate new span ID for this request
    span_id = generate_span_id()

    # Set trace context
    set_trace_context(trace_id, span_id, parent_span_id)

    # Add trace context to request for logging
    request["trace_id"] = trace_id
    request["span_id"] = span_id
    request["parent_span_id"] = parent_span_id

    # Start timing
    start_time = time.time()

    try:
        # Execute handler
        response = await handler(request)

        # Calculate duration
        duration_ms = (time.time() - start_time) * 1000

        # Add trace headers to response (W3C TraceContext)
        # Format: 00-{trace_id}-{span_id}-01 (01 = sampled)
        traceparent_value = f"00-{trace_id}-{span_id}-01"
        response.headers["traceparent"] = traceparent_value

        # Add trace metadata to response headers for debugging
        response.headers["X-Trace-ID"] = trace_id
        response.headers["X-Span-ID"] = span_id

        # Log span completion
        logger.debug(
            "Span completed: %s  trace=%s  span=%s  duration=%.1fms",
            request.path,
            trace_id[:16],
            span_id[:8],
            duration_ms,
            extra={
                "trace_id": trace_id,
                "span_id": span_id,
                "parent_span_id": parent_span_id,
                "duration_ms": duration_ms,
                "method": request.method,
                "path": request.path,
                "status": response.status,
            },
        )

        return response

    except Exception as e:
        duration_ms = (time.time() - start_time) * 1000

        # Log span error
        logger.error(
            "Span error: %s  trace=%s  span=%s  duration=%.1fms  error=%s",
            request.path,
            trace_id[:16],
            span_id[:8],
            duration_ms,
            str(e),
            extra={
                "trace_id": trace_id,
                "span_id": span_id,
                "parent_span_id": parent_span_id,
                "duration_ms": duration_ms,
                "method": request.method,
                "path": request.path,
                "error": str(e),
            },
            exc_info=True,
        )
        raise


def get_trace_context_for_logging() -> dict[str, str | None]:
    """Get trace context for logging."""
    return {
        "trace_id": get_trace_id(),
        "span_id": get_span_id(),
        "parent_span_id": _parent_span_id_var.get(),
    }


__all__ = [
    "generate_span_id",
    "generate_trace_id",
    "get_span_id",
    "get_trace_context_for_logging",
    "get_trace_id",
    "set_trace_context",
    "tracing_middleware",
]
