#!/usr/bin/env python3
"""Round-robin proxy in front of the frank-medium TP=1 replicas.

Why this exists: TP=2 (one vLLM process spanning both R9700s) needs RCCL
collectives, and RCCL's kernels require PCIe atomics on every participating
GPU. On this host the second R9700 sits behind chipset bridges that advertise
`AtomicOpsCap: Routing-` (see the runbook), so no collective can launch on it
and TP=2 is impossible. Two single-GPU replicas use both cards instead; this
proxy keeps the fleet contract of one endpoint on :7770.

The vLLM OpenAI API is stateless per request, so round-robin is sound. /health
returns 200 while at least one replica is up and names the ones that are not.
"""

from __future__ import annotations

import asyncio
import itertools
import os
from aiohttp import ClientSession, ClientTimeout, web

UPSTREAMS = [u.strip() for u in os.environ.get(
    "UPSTREAMS", "http://127.0.0.1:7771,http://127.0.0.1:7772").split(",") if u.strip()]
PORT = int(os.environ.get("PORT", "7770"))
# Hop-by-hop headers must not be forwarded (RFC 9110 §7.6.1) — forwarding
# transfer-encoding or content-length of a rewritten body breaks clients.
HOP_BY_HOP = {
    "host", "content-length", "transfer-encoding", "connection",
    "keep-alive", "proxy-authenticate", "proxy-authorization", "te",
    "trailers", "upgrade",
}

_cycle = itertools.cycle(range(len(UPSTREAMS)))


async def health(session: ClientSession) -> web.Response:
    results = await asyncio.gather(
        *(session.get(u + "/health", timeout=ClientTimeout(total=5)) for u in UPSTREAMS),
        return_exceptions=True,
    )
    lines, up = [], 0
    for url, res in zip(UPSTREAMS, results):
        if isinstance(res, Exception) or res.status != 200:
            lines.append(f"{url} DOWN")
        else:
            up += 1
            lines.append(f"{url} UP")
    body = "\n".join(lines) + "\n"
    return web.Response(status=200 if up else 503, text=body)


async def proxy(request: web.Request) -> web.StreamResponse:
    session: ClientSession = request.app["session"]
    if request.path == "/health":
        return await health(session)

    target = UPSTREAMS[next(_cycle)] + request.path_qs
    headers = {k: v for k, v in request.headers.items() if k.lower() not in HOP_BY_HOP}
    body = await request.read()

    try:
        upstream = await session.request(
            request.method, target, headers=headers, data=body or None,
            timeout=ClientTimeout(total=None, sock_connect=10, sock_read=3600),
        )
    except Exception as exc:  # upstream refused/unreachable
        return web.Response(status=502, text=f"proxy: {target} unreachable: {exc}\n")

    out = web.StreamResponse(status=upstream.status, reason=upstream.reason)
    for k, v in upstream.headers.items():
        if k.lower() not in HOP_BY_HOP:
            out.headers[k] = v
    await out.prepare(request)
    try:
        async for chunk in upstream.content.iter_any():
            await out.write(chunk)
    finally:
        upstream.release()
    await out.write_eof()
    return out


async def main() -> None:
    app = web.Application(client_max_size=1024 ** 3)
    app["session"] = ClientSession()
    app.router.add_route("*", "/{tail:.*}", proxy)
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, "0.0.0.0", PORT)
    await site.start()
    print(f"proxy: :{PORT} -> {', '.join(UPSTREAMS)}", flush=True)
    await asyncio.Event().wait()


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