"""
Correlation ID Middleware for FastAPI
Adds unique request ID for distributed tracing
"""
from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import Response
import uuid
import time
from typing import Callable

class CorrelationIdMiddleware(BaseHTTPMiddleware):
    """
    Middleware that adds correlation ID to all requests for distributed tracing
    """
    
    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        # Check for existing correlation ID in headers
        existing_id = (
            request.headers.get('x-correlation-id') or
            request.headers.get('x-request-id') or
            request.headers.get('traceparent')
        )
        
        # Generate new ID if not provided
        correlation_id = existing_id or self._generate_correlation_id()
        
        # Attach to request state
        request.state.correlation_id = correlation_id
        request.state.request_start_time = time.time()
        
        # Process request
        response = await call_next(request)
        
        # Add correlation ID to response headers
        response.headers['X-Correlation-ID'] = correlation_id
        response.headers['X-Request-ID'] = correlation_id
        
        return response
    
    @staticmethod
    def _generate_correlation_id() -> str:
        """
        Generate a unique correlation ID
        Format: timestamp-random (e.g., 1703423456789-a1b2c3d4)
        """
        timestamp = int(time.time() * 1000)
        random_part = uuid.uuid4().hex[:8]
        return f"{timestamp}-{random_part}"

def get_correlation_id(request: Request) -> str:
    """
    Get correlation ID from request (for use in services)
    
    Args:
        request: FastAPI request object
    
    Returns:
        Correlation ID string
    """
    return getattr(request.state, 'correlation_id', 'unknown')

def generate_correlation_id() -> str:
    """
    Generate a new correlation ID
    
    Returns:
        Correlation ID string
    """
    return CorrelationIdMiddleware._generate_correlation_id()