import time
from typing import Callable, Dict, Optional, Union

from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import Response

from app.core.logger import logger


class RequestLoggingMiddleware(BaseHTTPMiddleware):
    MAX_BODY_LOG_SIZE = 1_000_000

    def __init__(
        self,
        app,
        log_request_body: bool = True,
        log_query_params: bool = True,
        exclude_paths: Optional[list[str]] = None,
    ) -> None:
        super().__init__(app)
        self.log_request_body = log_request_body
        self.log_query_params = log_query_params
        self.exclude_paths = set(exclude_paths or ["/health", "/metrics"])

    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        if request.url.path in self.exclude_paths:
            return await call_next(request)

        request_body: Optional[str] = None
        query_params: Optional[dict[str, str]] = None

        if self.log_request_body:
            request_body = await self._get_request_body(request)
        if self.log_query_params:
            params = dict(request.query_params)
            query_params = params if params else None

        start_time = time.perf_counter()

        try:
            response = await call_next(request)
        except Exception:
            duration_ms = round((time.perf_counter() - start_time) * 1000, 2)
            logger.exception(
                "request_failed",
                method=request.method,
                path=request.url.path,
                duration_ms=duration_ms,
                request_body=request_body,
                query_params=query_params,
                client_ip=self._get_client_ip(request),
            )
            raise

        duration_ms = round((time.perf_counter() - start_time) * 1000, 2)

        log_kwargs: Dict[str, Union[str, int, float, dict, None]] = {
            "method": request.method,
            "path": request.url.path,
            "status_code": response.status_code,
            "duration_ms": duration_ms,
            "client_ip": self._get_client_ip(request),
        }

        if request_body:
            log_kwargs["request_body"] = request_body
        if query_params:
            log_kwargs["query_params"] = query_params

        if response.status_code >= 500:
            logger.error("request_completed", **log_kwargs)
        elif response.status_code >= 400:
            logger.warning("request_completed", **log_kwargs)
        else:
            logger.info("request_completed", **log_kwargs)

        return response

    @staticmethod
    def _get_client_ip(request: Request) -> str:
        """Extract client IP from request, considering proxies."""
        forwarded = request.headers.get("X-Forwarded-For")
        if forwarded:
            return forwarded.split(",")[0].strip()
        if request.client:
            return request.client.host
        return "unknown"

    async def _get_request_body(self, request: Request) -> Optional[str]:
        """Read and return request body for logging."""
        if request.method not in {"POST", "PUT", "PATCH"}:
            return None

        content_type = request.headers.get("content-type", "")
        if content_type.startswith("multipart/form-data"):
            return "[multipart/form-data]"

        try:
            body_bytes = await request.body()

            async def receive() -> Dict[str, Union[str, bytes]]:
                return {"type": "http.request", "body": body_bytes}

            request._receive = receive

            if len(body_bytes) > self.MAX_BODY_LOG_SIZE:
                return "[body too large]"

            return body_bytes.decode("utf-8") if body_bytes else None
        except (UnicodeDecodeError, RuntimeError):
            return "[unable to read body]"
