#!/usr/bin/env python3
"""Lightweight round-robin reverse proxy for vLLM data-parallel child servers.

vLLM's --data-parallel-multi-port-external-lb only exposes aggregated /health on
a supervisor port; it does not forward /v1/models, /v1/completions, etc.  This
proxy sits in front of the child API servers and round-robins requests across
them.

Usage:
    python3 dp_proxy.py --port 30040 --backends 30030,30031,30032,30033
"""
import argparse
import asyncio
import itertools
from contextlib import asynccontextmanager

import httpx
from fastapi import FastAPI, Request, Response
from fastapi.responses import StreamingResponse
from uvicorn import Config, Server


HOP_BY_HOP_HEADERS = {
    "connection",
    "keep-alive",
    "proxy-authenticate",
    "proxy-authorization",
    "te",
    "trailers",
    "transfer-encoding",
    "upgrade",
    "content-encoding",
    "content-length",
}


class DPProxy:
    def __init__(self, backends: list[int], timeout: float = 60.0):
        self.backend_urls = [f"http://127.0.0.1:{port}" for port in backends]
        self.cycle = itertools.cycle(self.backend_urls)
        self.lock = asyncio.Lock()
        self.timeout = timeout
        self.client = httpx.AsyncClient(
            timeout=httpx.Timeout(timeout),
            limits=httpx.Limits(max_connections=100, max_keepalive_connections=20),
        )

    async def close(self):
        await self.client.aclose()

    async def _next_backend(self) -> str:
        async with self.lock:
            return next(self.cycle)

    async def _is_healthy(self, base_url: str) -> bool:
        try:
            r = await self.client.get(f"{base_url}/health", timeout=2.0)
            return r.status_code == 200
        except Exception:
            return False

    async def health(self) -> Response:
        results = await asyncio.gather(
            *(self._is_healthy(url) for url in self.backend_urls)
        )
        if all(results):
            return Response(status_code=200)
        return Response(status_code=503)

    async def ready(self) -> Response:
        # Bench clients use /v1/models as the ready signal.  We forward that
        # to a healthy backend; this endpoint just checks /health.
        return await self.health()

    async def proxy(self, request: Request) -> Response:
        path = request.url.path
        if request.url.query:
            path = f"{path}?{request.url.query}"

        body = await request.body()
        headers = dict(request.headers)
        headers.pop("host", None)

        tried = set()
        while len(tried) < len(self.backend_urls):
            base_url = await self._next_backend()
            if base_url in tried:
                # Avoid infinite loop when only one backend remains.
                break
            tried.add(base_url)

            target = f"{base_url}{path}"
            try:
                backend_req = self.client.build_request(
                    method=request.method,
                    url=target,
                    headers=headers,
                    content=body,
                )
                backend_resp = await self.client.send(
                    backend_req, follow_redirects=False, stream=True
                )
            except Exception:
                continue

            if backend_resp.status_code >= 500:
                await backend_resp.aclose()
                continue

            response_headers = {
                k: v
                for k, v in backend_resp.headers.items()
                if k.lower() not in HOP_BY_HOP_HEADERS
            }
            return StreamingResponse(
                backend_resp.aiter_bytes(),
                status_code=backend_resp.status_code,
                headers=response_headers,
                background=None,
            )

        return Response(
            content=b"All DP backends are unavailable",
            status_code=503,
        )


def make_app(proxy: DPProxy) -> FastAPI:
    app = FastAPI()

    @app.get("/health", include_in_schema=False)
    async def health():
        return await proxy.health()

    @app.get("/ready", include_in_schema=False)
    @app.get("/readyz", include_in_schema=False)
    async def ready():
        return await proxy.ready()

    @app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"])
    async def catch_all(request: Request):
        return await proxy.proxy(request)

    return app


async def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", type=int, required=True)
    parser.add_argument(
        "--backends",
        type=lambda s: [int(x.strip()) for x in s.split(",")],
        required=True,
        help="Comma-separated list of child vLLM server ports",
    )
    parser.add_argument("--timeout", type=float, default=300.0)
    parser.add_argument("--log-level", default="warning")
    args = parser.parse_args()

    proxy = DPProxy(args.backends, timeout=args.timeout)
    app = make_app(proxy)

    config = Config(app, host="0.0.0.0", port=args.port, log_level=args.log_level)
    server = Server(config)

    try:
        await server.serve()
    finally:
        await proxy.close()


if __name__ == "__main__":
    asyncio.run(main())