from __future__ import annotations

import time
import uuid
from typing import Callable

from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.types import ASGIApp


class RequestLoggingMiddleware(BaseHTTPMiddleware):
    """
    Middleware for logging requests and responses.
    """

    def __init__(self, app: ASGIApp, logger: Callable = print):
        super().__init__(app)
        self.logger = logger

    async def dispatch(self, request: Request, call_next):
        request_id = str(uuid.uuid4())
        request.state.request_id = request_id

        # Log request
        start_time = time.time()
        self.logger(f"Request {request_id} started: {request.method} {request.url.path}")

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

            # Log response
            process_time = time.time() - start_time
            self.logger(
                f"Request {request_id} completed: {response.status_code} ({process_time:.4f}s)"
            )

            # Add custom headers
            response.headers["X-Request-ID"] = request_id
            response.headers["X-Process-Time"] = str(process_time)

            return response
        except Exception as e:
            # Log exception
            process_time = time.time() - start_time
            self.logger(f"Request {request_id} failed: {e!s} ({process_time:.4f}s)")
            raise