import json
import time
import logging
logger = logging.getLogger(__name__)
from fastapi import FastAPI, Request
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import StreamingResponse

class LoggingMiddleware(BaseHTTPMiddleware):
    """HTTP请求/响应日志中间件."""

    async def dispatch(self, request: Request, call_next):
        # 记录请求开始时间
        start_time = time.time()

        # 获取请求信息
        method = request.method
        url = str(request.url)
        client_host = request.client.host if request.client else "unknown"

        # 读取请求体
        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