import json
import time
from typing import Callable

from fastapi import APIRouter, Request, Response
from fastapi.responses import JSONResponse
from fastapi.routing import APIRoute

from config.setting import APP_VERSION


class TimedRoute(APIRoute):
    def get_route_handler(self) -> Callable:
        original_route_handler = super().get_route_handler()

        async def custom_route_handler(request: Request) -> Response:
            start = time.perf_counter()
            response: Response = await original_route_handler(request)
            duration = time.perf_counter() - start

            # JSON 응답만 처리
            if getattr(response, "media_type", None) == "application/json":
                body_bytes = b""

                # 1) 스트리밍 응답이면 모두 수집
                body_iter = getattr(response, "body_iterator", None)
                if body_iter is not None:
                    async for chunk in body_iter:
                        body_bytes += chunk
                else:
                    # 일반 응답
                    body_bytes = getattr(response, "body", b"")

                # 2) 파싱 시도
                try:
                    text = body_bytes.decode("utf-8") if body_bytes else ""
                    payload = json.loads(text) if text else None
                except Exception:
                    # 파싱 실패: 바디/길이 불일치 위험 없도록 원본 그대로 반환
                    return response

                # 3) dict일 때만 필드 주입
                if isinstance(payload, dict):
                    payload.setdefault("app_version", APP_VERSION)
                    payload.setdefault("process_time", f"{duration:.4f}")

                # 4) 새 JSONResponse로 재구성 (Content-Length 자동 일치)
                #    원본 상태코드/헤더/미디어타입 유지
                new_headers = dict(response.headers)
                # 압축/길이 관련 헤더는 제거(재계산되도록)
                for h in ("content-length", "Content-Length", "content-encoding", "Content-Encoding"):
                    new_headers.pop(h, None)

                return JSONResponse(
                    content=payload,
                    status_code=response.status_code,
                    headers=new_headers,
                    media_type=response.media_type,
                )

            # JSON 아니라면 바디 불문, 원본 그대로
            return response

        return custom_route_handler


class CustomAPIRouter(APIRouter):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.route_class = TimedRoute