"""Request correlation utilities for Robin AI.

Provides utilities for maintaining correlation IDs across requests and services,
enabling distributed tracing and log correlation throughout the system.
"""

import uuid
from contextvars import ContextVar
from typing import Any

from opentelemetry import trace
from opentelemetry.trace import format_span_id, format_trace_id

# Updated import to use robin_commons.log
from robin_commons.log import logger

# Context variables for request correlation
_correlation_id: ContextVar[str | None] = ContextVar("correlation_id", default=None)
_request_id: ContextVar[str | None] = ContextVar("request_id", default=None)
_user_id: ContextVar[str | None] = ContextVar("user_id", default=None)
_trace_id: ContextVar[str | None] = ContextVar("trace_id", default=None)
_span_id: ContextVar[str | None] = ContextVar("span_id", default=None)


def generate_correlation_id() -> str:
    """Generate a new correlation ID."""
    return str(uuid.uuid4())


def set_correlation_id(correlation_id: str) -> None:
    """Set the correlation ID for the current context.

    Args:
        correlation_id: The correlation ID to set
    """
    _correlation_id.set(correlation_id)
    logger.contextualize(correlation_id=correlation_id)


def get_correlation_id() -> str | None:
    """Get the current correlation ID.

    Returns:
        The current correlation ID or None if not set
    """
    return _correlation_id.get()


def set_request_id(request_id: str) -> None:
    """Set the request ID for the current context.

    Args:
        request_id: The request ID to set
    """
    _request_id.set(request_id)
    logger.contextualize(request_id=request_id)


def get_request_id() -> str | None:
    """Get the current request ID.

    Returns:
        The current request ID or None if not set
    """
    return _request_id.get()


def set_user_id(user_id: str) -> None:
    """Set the user ID for the current context.

    Args:
        user_id: The user ID to set
    """
    _user_id.set(user_id)
    logger.contextualize(user_id=user_id)


def get_user_id() -> str | None:
    """Get the current user ID.

    Returns:
        The current user ID or None if not set
    """
    return _user_id.get()


def set_trace_context(trace_id: str | None = None, span_id: str | None = None) -> None:
    """Set the OpenTelemetry trace context for the current request.

    Args:
        trace_id: The trace ID to set (will be extracted from current span if None)
        span_id: The span ID to set (will be extracted from current span if None)
    """
    # Extract from current span if not provided
    if trace_id is None or span_id is None:
        current_span = trace.get_current_span()
        if current_span and current_span.is_recording():
            span_context = current_span.get_span_context()
            if trace_id is None:
                trace_id = format_trace_id(span_context.trace_id)
            if span_id is None:
                span_id = format_span_id(span_context.span_id)

    # Set context variables
    if trace_id:
        _trace_id.set(trace_id)
    if span_id:
        _span_id.set(span_id)

    # Contextualize logger with available IDs
    context = {}
    if trace_id:
        context["trace_id"] = trace_id
    if span_id:
        context["span_id"] = span_id
    if context:
        logger.contextualize(**context)


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

    Returns:
        The current trace ID or None if not set
    """
    return _trace_id.get()


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

    Returns:
        The current span ID or None if not set
    """
    return _span_id.get()


def get_correlation_context() -> dict[str, Any]:
    """Get all correlation context as a dictionary.

    Returns:
        Dictionary containing all correlation context
    """
    context = {}

    if correlation_id := get_correlation_id():
        context["correlation_id"] = correlation_id

    if request_id := get_request_id():
        context["request_id"] = request_id

    if user_id := get_user_id():
        context["user_id"] = user_id

    if trace_id := get_trace_id():
        context["trace_id"] = trace_id

    if span_id := get_span_id():
        context["span_id"] = span_id

    return context


def set_correlation_context(context: dict[str, Any]) -> None:
    """Set correlation context from a dictionary.

    Args:
        context: Dictionary containing correlation context
    """
    if correlation_id := context.get("correlation_id"):
        set_correlation_id(correlation_id)

    if request_id := context.get("request_id"):
        set_request_id(request_id)

    if user_id := context.get("user_id"):
        set_user_id(user_id)

    if trace_id := context.get("trace_id"):
        if span_id := context.get("span_id"):
            set_trace_context(trace_id, span_id)
        else:
            set_trace_context(trace_id=trace_id)


def clear_correlation_context() -> None:
    """Clear all correlation context."""
    _correlation_id.set(None)
    _request_id.set(None)
    _user_id.set(None)
    _trace_id.set(None)
    _span_id.set(None)


def ensure_correlation_id() -> str:
    """Ensure a correlation ID exists, generating one if necessary.

    Returns:
        The current or newly generated correlation ID
    """
    correlation_id = get_correlation_id()
    if not correlation_id:
        correlation_id = generate_correlation_id()
        set_correlation_id(correlation_id)
    return correlation_id


def get_correlation_headers() -> dict[str, str]:
    """Get correlation context as HTTP headers for propagation.

    Returns:
        Dictionary of HTTP headers for correlation propagation
    """
    headers = {}

    if correlation_id := get_correlation_id():
        headers["X-Correlation-ID"] = correlation_id

    if request_id := get_request_id():
        headers["X-Request-ID"] = request_id

    if user_id := get_user_id():
        headers["X-User-ID"] = user_id

    return headers


def set_correlation_from_headers(headers: dict[str, str]) -> None:
    """Set correlation context from HTTP headers.

    Args:
        headers: Dictionary of HTTP headers (case-insensitive)
    """
    # Convert to lowercase for case-insensitive lookup
    lower_headers = {k.lower(): v for k, v in headers.items()}

    if correlation_id := lower_headers.get("x-correlation-id"):
        set_correlation_id(correlation_id)

    if request_id := lower_headers.get("x-request-id"):
        set_request_id(request_id)

    if user_id := lower_headers.get("x-user-id"):
        set_user_id(user_id)


def add_span_attributes(span: trace.Span | None = None) -> None:
    """Add correlation context as span attributes.

    Args:
        span: The span to add attributes to (uses current span if None)
    """
    if span is None:
        span = trace.get_current_span()

    if not span or not span.is_recording():
        return

    context = get_correlation_context()
    for key, value in context.items():
        span.set_attribute(f"correlation.{key}", value)


def log_correlation_context(message: str = "Correlation context") -> None:
    """Log the current correlation context for debugging.

    Args:
        message: Log message prefix
    """
    context = get_correlation_context()
    if context:
        logger.debug(f"{message}: {context}")
    else:
        logger.debug(f"{message}: No correlation context set")
