#!/usr/bin/env python3
"""HTTP front door for the frank-medium single-GPU replicas.

Why two replicas instead of TP=2: one vLLM process spanning both R9700s needs
RCCL collectives, and RCCL's kernels require PCIe atomics on every
participating GPU. The second card sits behind chipset bridges advertising
`AtomicOpsCap: Routing-`, which no software can enable (see the runbook), so
no collective can launch on it. Two single-GPU replicas use both cards; this
service keeps the fleet contract of one endpoint on :7770, streaming included.

Routing decisions live in routing.py — spread single-turn requests, pin
continuations to the replica whose prefix cache holds their context. This file
only moves bytes: parse enough of the request to route it, forward it verbatim,
stream the answer back, and report replica health on /health.
"""

from __future__ import annotations

import asyncio
import os
from contextlib import suppress
from aiohttp import ClientSession, ClientTimeout, web

import routing

PORT = int(os.environ.get("PORT", "7770"))
HEALTH_TTL = float(os.environ.get("HEALTH_TTL", "10"))
# 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",
}

UPSTREAMS = routing.UPSTREAMS
# Replica health, refreshed by the watcher. Module state rather than aiohttp app
# state: the application is frozen once the site is up.
_UP: list[bool] = [True] * 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")
    return web.Response(status=200 if up else 503, text="\n".join(lines) + "\n")


async def watch_health() -> None:
    """Refresh the replica up/down table so routing never targets a dead one.

    One HTTP request per replica per interval, in exchange for never handing a
    client a 502 after a restart.
    """
    global _UP
    session = ClientSession()
    try:
        while True:
            res = await asyncio.gather(
                *(session.get(u + "/health", timeout=ClientTimeout(total=5)) for u in UPSTREAMS),
                return_exceptions=True,
            )
            _UP = [not isinstance(r, Exception) and r.status == 200 for r in res]
            await asyncio.sleep(HEALTH_TTL)
    finally:
        await session.close()


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

    body = await request.read()
    headers = {k: v for k, v in request.headers.items() if k.lower() not in HOP_BY_HOP}

    # Fail over before anything is written to the client. A replica that
    # restarts is invisible to the health watcher for up to HEALTH_TTL, and a
    # deploy restarts both replicas — without this retry a 10 s window of
    # requests would 502 on a service that has a live card the whole time.
    # vLLM's OpenAI endpoints are stateless per request, so re-sending a
    # request that never reached an engine is safe; a request that failed
    # mid-stream is not retried (that path is below).
    tried: set[int] = set()
    upstream = None
    last_error = "no healthy frank-medium replica"
    while len(tried) < len(UPSTREAMS):
        idx = routing.route(body, request.remote, _UP)
        if idx < 0 or idx in tried:
            break
        tried.add(idx)
        target = UPSTREAMS[idx] + request.path_qs
        try:
            upstream = await session.request(
                request.method, target, headers=headers, data=body or None,
                timeout=ClientTimeout(total=None, sock_connect=10, sock_read=3600),
            )
            break
        except Exception as exc:
            # Demote at once: the watcher will confirm, but every request until
            # then should skip this replica rather than pay the same timeout.
            # Logged because a silent failover hides a card that is down.
            _UP[idx] = False
            last_error = f"{target} unreachable: {exc}"
            print(f"proxy: failover — {target} unreachable ({exc}); retrying "
                  f"another replica", flush=True)
    if upstream is None:
        # 503 when no replica was even reachable through the routing table,
        # 502 when replicas were tried and all of them failed.
        status = 502 if tried else 503
        return web.Response(status=status, text=f"proxy: {last_error}\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)
    except Exception:
        # The client went away mid-stream (an aborted generation). Nothing to
        # report: the engine sees its own cancellation.
        pass
    finally:
        upstream.release()
    with suppress(Exception):
        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()
    await web.TCPSite(runner, "0.0.0.0", PORT).start()
    print(f"proxy: :{PORT} -> {', '.join(UPSTREAMS)} routing={routing.ROUTING} "
          f"key={routing.AFFINITY_BYTES}B known_max={routing.KNOWN_MAX}", flush=True)
    watcher = asyncio.create_task(watch_health())
    with suppress(asyncio.CancelledError):
        await asyncio.gather(watcher)


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