# !/usr/bin/env python
# -*-coding:utf-8 -*-

"""
# Time       ：2024/1/12 10:27
# Author     ：Maxwell
# Description：
"""
import time
import traceback
from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware
from infoman.logger import logger
from infoman.utils.http.info import ClientInfoExtractor
from starlette.responses import StreamingResponse, FileResponse


class LoggingMiddleware(BaseHTTPMiddleware):

    @staticmethod
    def format_size(size_bytes: int) -> str:
        if size_bytes < 1024:
            return f"{size_bytes}B"
        elif size_bytes < 1024 * 1024:
            return f"{size_bytes / 1024:.2f}KB"
        else:
            return f"{size_bytes / (1024 * 1024):.2f}MB"

    async def dispatch(self, request: Request, call_next) -> Response:
        start_time = time.monotonic()
        client_ip = ClientInfoExtractor.ip_address(request=request)
        path = request.url.path[:200]

        response = await call_next(request)
        elapsed_ms = int((time.monotonic() - start_time) * 1000)

        content_length = response.headers.get("content-length")
        if content_length:
            response_size = int(content_length)
            size_str = self.format_size(response_size)
        else:
            size_str = "unknown"

        logger.info(
            f"Req: ip={client_ip}, elapsed_ms={elapsed_ms}, "
            f"path={path}, status={response.status_code}, "
            f"size={size_str}"
        )
        return response


