import asyncio, subprocess, sys, time, statistics, httpx, fastapi, starlette, uvicorn, platform

def start(port, gzip):
    code = f"import uvicorn,app;uvicorn.run(app.make_app(gzip={gzip}),port={port},log_level='warning')"
    p = subprocess.Popen([sys.executable, "-c", code])
    for _ in range(100):
        try:
            httpx.get(f"http://127.0.0.1:{port}/openapi.json"); return p
        except Exception: time.sleep(0.1)
    raise RuntimeError

async def one(client, path):
    t0 = time.perf_counter(); first = None; src = None; n = 0; chunks = 0
    async with client.stream("GET", path) as r:
        async for chunk in r.aiter_raw():
            now = time.perf_counter() - t0
            if first is None: first = now
            chunks += 1
            if src is None and b"event: sources" in chunk: src = now
            if b"event: token" in chunk and n == 0: n = now
        total = time.perf_counter() - t0
    return dict(ttfb=first, src=src, tok=n or None, total=total, chunks=chunks)

def med(xs): return statistics.median(xs)*1000
async def run(port, path, conc, reps):
    async with httpx.AsyncClient(base_url=f"http://127.0.0.1:{port}", timeout=120) as c:
        await one(c, path)  # warmup
        out = []
        for _ in range(reps):
            out += await asyncio.gather(*[one(c, path) for _ in range(conc)])
    return out

def row(name, out):
    f = lambda k: f"{med([o[k] for o in out if o[k] is not None]):.0f}" if any(o[k] is not None for o in out) else "-"
    mx = max(o["ttfb"] for o in out)*1000
    print(f"| {name} | {f('ttfb')} | {f('src')} | {f('tok')} | {f('total')} | {mx:.0f} |")

async def main():
    print(f"fastapi {fastapi.__version__}, starlette {starlette.__version__}, uvicorn {uvicorn.__version__}, httpx {httpx.__version__}, python {platform.python_version()}")
    print("simulated: retrieval 50 ms, LLM first token 400 ms, 60 tokens x 20 ms\n")
    print("| Case | median first byte ms | median sources ms | median first token ms | median total ms | worst first byte ms |")
    print("|---|---|---|---|---|---|")
    p = start(8101, False); p2 = start(8102, True)
    try:
        for name, port, path, conc in [
            ("non-streaming JSON, 1 client", 8101, "/plain", 1),
            ("SSE sources-first, 1 client", 8101, "/stream", 1),
            ("SSE sources-first, 20 concurrent", 8101, "/stream", 20),
            ("SSE + blocking time.sleep, 5 concurrent", 8101, "/stream-blocking", 5),
            ("SSE sources-first, 5 concurrent (control)", 8101, "/stream", 5),
            ("SSE + GZipMiddleware, 1 client", 8102, "/stream", 1)]:
            row(name, await run(port, path, conc, 5 if conc == 1 else 3) if conc>1 else await run(port, path, 1, 15))
        # gzip: request with accept-encoding
        async with httpx.AsyncClient(base_url="http://127.0.0.1:8102", headers={"Accept-Encoding":"gzip"}, timeout=120) as c:
            await one(c, "/stream"); outs = [await one(c, "/stream") for _ in range(10)]
        row("SSE + GZipMiddleware, client sends Accept-Encoding: gzip", outs)
        print("\nchunks received per response (1 client, no gzip):", (await run(8101, "/stream", 1, 1))[0]["chunks"], "of 62 events")
        print("chunks with gzip + Accept-Encoding:", outs[0]["chunks"])
    finally:
        p.terminate(); p2.terminate()
asyncio.run(main())