import json
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse
from starlette.exceptions import HTTPException
from starlette.requests import Request
from common.core.config import settings
from common.utils.utils import SQLBotLogUtil
class ResponseMiddleware(BaseHTTPMiddleware):
    def __init__(self, app):
        super().__init__(app)

    async def dispatch(self, request, call_next):
        response = await call_next(request)
        
        if isinstance(response, JSONResponse) or request.url.path == f"{settings.API_V1_STR}/openapi.json":
            return response
        if response.status_code != 200:
            return response
        if response.headers.get("content-type") == "application/json":
            try:
                body = b""
                async for chunk in response.body_iterator:
                    body += chunk
                
                raw_data = json.loads(body.decode())
                
                if isinstance(raw_data, dict) and all(k in raw_data for k in ["code", "data", "msg"]):
                    return JSONResponse(
                        content=raw_data,
                        status_code=response.status_code,
                        headers={
                            k: v for k, v in response.headers.items()
                            if k.lower() not in ("content-length", "content-type")
                        }
                    )
                
                wrapped_data = {
                    "code": 0,
                    "data": raw_data,
                    "msg": None
                }
                
                return JSONResponse(
                    content=wrapped_data,
                    status_code=response.status_code,
                    headers={
                        k: v for k, v in response.headers.items()
                        if k.lower() not in ("content-length", "content-type")
                    }
                )
            except Exception as e:
                SQLBotLogUtil.error(f"Response processing error: {str(e)}", exc_info=True)
                return JSONResponse(
                    status_code=500,
                    content=str(e),
                    headers={
                        k: v for k, v in response.headers.items()
                        if k.lower() not in ("content-length", "content-type")
                    }
                )
                
        return response


class exception_handler():
    @staticmethod
    async def http_exception_handler(request: Request, exc: HTTPException):
        SQLBotLogUtil.error(f"HTTP Exception: {exc.detail}", exc_info=True)
        return JSONResponse(
            status_code=exc.status_code,
            content=exc.detail,
            headers={"Access-Control-Allow-Origin": "*"}
        )


    @staticmethod
    async def global_exception_handler(request: Request, exc: Exception):
        SQLBotLogUtil.error(f"Unhandled Exception: {str(exc)}", exc_info=True)
        return JSONResponse(
            status_code=500,
            content=str(exc),
            headers={"Access-Control-Allow-Origin": "*"}
        )

