from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import Response, JSONResponse
import json
import os

from app.models.api_response import APIResponseModel
from app.utils.response import APIResponse

class ResponseWrapperMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        # 检查是否需要启用标准响应包装
        use_standard_response = os.environ.get("USE_STANDARD_RESPONSE") == "1"
        
        # 不包装的情况：健康检查端点或禁用了标准响应
        health_endpoints = ["/health", "/"]
        if request.url.path in health_endpoints or not use_standard_response:
            return await call_next(request)
        
        response = await call_next(request)
        
        # 如果响应已经被包装或者是文件下载等特殊响应，则不再处理
        if (
            isinstance(response, Response) 
            and not isinstance(response, JSONResponse)
            or response.headers.get("content-type") != "application/json"
        ):
            return response
        
        # 处理JSON响应
        try:
            response_body = [section async for section in response.body_iterator]
            response.body_iterator = None
            
            if len(response_body) > 0:
                body = response_body[0].decode()
                json_body = json.loads(body)
                
                # 检查是否已经是标准格式
                if isinstance(json_body, dict) and "code" in json_body and "data" in json_body and "message" in json_body:
                    # 已经是标准格式，直接返回
                    new_response = Response(
                        content=body,
                        status_code=response.status_code,
                        headers=dict(response.headers),
                        media_type=response.media_type
                    )
                    return new_response
                
                # 包装响应
                api_response = APIResponseModel(
                    code=200,
                    message="操作成功",
                    data=json_body
                )
                
                return JSONResponse(
                    content=api_response.dict(),
                    status_code=response.status_code,
                    headers=dict(response.headers),
                )
            
            return response
        except Exception as e:
            # 如果解析失败，返回原始响应
            return response