import json
import time
import logging

logger = logging.getLogger(__name__)
from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import StreamingResponse
from config import settings


class LoggingMiddleware(BaseHTTPMiddleware):

    # 日志记录路径前缀
    async def dispatch(self, request: Request, call_next):
        # 获取请求信息
        method = request.method
        url = str(request.url)
        client_host = request.client.host if request.client else "unknown"

        # 判断是否需要记录日志
        should_log = request.url.path.startswith(settings.url_prefix)

        if should_log:
            # 记录请求开始时间
            start_time = time.time()

            # 读取请求体
            request_body = None
            if method in ("POST", "PUT", "PATCH"):
                try:
                    body = await request.body()
                    if body:
                        request_body = body.decode("utf-8", errors="replace")
                        # 尝试格式化为JSON
                        try:
                            request_body = json.loads(request_body)
                        except json.JSONDecodeError:
                            pass  # 保持原始字符串
                except Exception:
                    pass

            # 记录请求日志
            logger.info(f"→ HTTP Request | {method} {url} | Client: {client_host}")
            if request_body:
                logger.info(f"  Request Body: {request_body}")

            # 处理请求
            response = await call_next(request)

            # 计算处理时间
            process_time = time.time() - start_time

            # 读取响应体（仅处理非流式响应）
            response_body = None
            if not isinstance(response, StreamingResponse):
                try:
                    response_body_bytes = b""
                    async for chunk in response.body_iterator:
                        response_body_bytes += chunk

                    # 重新构建响应
                    response_body = response_body_bytes.decode("utf-8", errors="replace")
                    # 尝试格式化为JSON
                    try:
                        response_body = json.loads(response_body)
                    except json.JSONDecodeError:
                        pass  # 保持原始字符串

                    # 重建响应对象
                    response = StreamingResponse(
                        iter([response_body_bytes]),
                        status_code=response.status_code,
                        headers=dict(response.headers),
                        media_type=response.media_type,
                    )
                except Exception:
                    pass

            # 记录响应日志
            status_code = response.status_code
            logger.info(f"← HTTP Response | {method} {url} | Status: {status_code} | Time: {process_time:.3f}s")
            if response_body:
                logger.info(f"  Response Body: {response_body}")

            return response
        else:
            # 不需要记录日志的请求，直接处理
            return await call_next(request)