from __future__ import annotations

from typing import Callable

from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware


class V1HeadersMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next: Callable):
        response = await call_next(request)
        # Only tag /api/v1 responses
        if request.url.path.startswith("/api/v1"):
            # Reuse request_id set by RequestContextMiddleware
            req_id = getattr(request.state, "request_id", "")
            if req_id:
                response.headers.setdefault("X-Foundry-Request-Id", req_id)
            # Best-effort rate-limit headers from dependency
            limit = getattr(request.state, "x_ratelimit_limit", None)
            remaining = getattr(request.state, "x_ratelimit_remaining", None)
            if limit is not None:
                response.headers["X-RateLimit-Limit"] = str(limit)
            if remaining is not None:
                response.headers["X-RateLimit-Remaining"] = str(remaining)
        return response


def install_v1_middleware(app):
    app.add_middleware(V1HeadersMiddleware)
