import asyncio
import sys
import time

import httpx
import openai

BASE_URL = sys.argv[1] if len(sys.argv) > 1 else "http://localhost:30000"
PATTERN = sys.argv[2] if len(sys.argv) > 2 else "1v3"

async def send(client, dp_rank, req_id):
    try:
        t0 = time.time()
        await client.chat.completions.create(
            model="GLM-5-FP8-DEBUG",
            messages=[
                {
                    "role": "user",
                    "content": f"Write an essay about request #{req_id}, be creative.",
                }
            ],
            max_tokens=128,
            extra_body={"routed_dp_rank": dp_rank},
        )
        return True, time.time() - t0
    except Exception as e:
        return False, str(e)[:100]

async def main():
    client = openai.AsyncOpenAI(base_url=f"{BASE_URL}/v1", api_key="none")

    async with httpx.AsyncClient() as hc:
        r = await hc.get(f"{BASE_URL}/health_generate", timeout=5)
        print(f"Health: {r.status_code}")

    print("--- warmup ---")
    print("  warmup:", await send(client, 0, 0))

    dp0, dp1 = map(int, PATTERN.split("v"))
    tasks = []
    rid = 100
    for _ in range(dp0):
        tasks.append(send(client, 0, rid))
        rid += 1
    for _ in range(dp1):
        tasks.append(send(client, 1, rid))
        rid += 1

    results = await asyncio.gather(*tasks, return_exceptions=True)
    ok = sum(1 for r in results if r is not None and r[0])
    print(f"Total: {ok}/{len(results)} ok")
    for r in results:
        print(" ", r)

asyncio.run(main())