#!/usr/bin/env python3
"""GLM host showdown sweep — stdlib-only port of glm-host-sweep.mjs so it can run
on hosts without node (e.g. kruti-ext1, colocated with the gateway — the blog's
measurement vantage). Same methodology: GLM-5.2, long-output streaming, N runs
per host round-robin, median latency/TTFT/tok/s, latency normalized to ~700 tok.

  BR_KEY=... python3 glm_host_sweep.py [--runs 20] [--dry] [--hosts fireworks,baseten]
  UPSTREAM_KEY=... adds a per-request BYOK key (use with a single --hosts entry).
"""
import argparse, json, os, statistics, sys, time, urllib.request

API = os.environ.get('BR_API', 'https://api.bharatrouter.com')
KEY = os.environ.get('BR_KEY')
if not KEY:
    sys.exit('Set BR_KEY')

# INR/Mtok, src/catalog.ts glm-5.2 (updated 2026-07-02); Fireworks list $1.40/$4.40 @ ₹96.
ALL_HOSTS = [
    ('Baseten',    {'model': 'glm-5.2', 'provider': 'baseten'},    (134, 422)),
    ('OpenRouter', {'model': 'glm-5.2', 'provider': 'openrouter'}, (89, 288)),
    ('Zhipu',      {'model': 'glm-5.2', 'provider': 'zhipu'},      (134, 422)),
    ('Fireworks',  {'model': 'fireworks/accounts/fireworks/models/glm-5p2'}, (134, 422)),
]

PROMPT = ('Explain, for a backend engineer new to distributed systems, how the Raft consensus '
          'algorithm works: leader election, log replication, and safety guarantees. Use concrete '
          'examples with a 5-node cluster. Aim for roughly 700 words. Plain prose, no headings.')


def one(body, inr):
    t0 = time.monotonic()
    ttft = usage = err = None
    payload = dict(body, stream=True, stream_options={'include_usage': True}, max_tokens=900,
                   messages=[{'role': 'user', 'content': PROMPT}])
    up = os.environ.get('UPSTREAM_KEY')
    if up:
        payload['upstream_key'] = up
    req = urllib.request.Request(f'{API}/v1/chat/completions', data=json.dumps(payload).encode(),
                                 headers={'Authorization': f'Bearer {KEY}',
                                          'Content-Type': 'application/json',
                                          'User-Agent': 'br-bench/glm-host-sweep-py'})
    try:
        with urllib.request.urlopen(req, timeout=120) as res:
            buf = b''
            while True:
                chunk = res.read(4096)
                if not chunk:
                    break
                buf += chunk
                while b'\n' in buf:
                    line, buf = buf.split(b'\n', 1)
                    line = line.strip()
                    if not line.startswith(b'data: ') or line == b'data: [DONE]':
                        continue
                    c = json.loads(line[6:])
                    d = (c.get('choices') or [{}])[0].get('delta') or {}
                    if ttft is None and (d.get('content') or d.get('reasoning_content')):
                        ttft = time.monotonic() - t0
                    if c.get('usage'):
                        usage = c['usage']
    except Exception as e:  # noqa: BLE001 — record any failure as an errored run
        err = str(e)[:200]
    total = time.monotonic() - t0
    out = (usage or {}).get('completion_tokens', 0)
    pin = (usage or {}).get('prompt_tokens', 0)
    return {'err': err, 'ttftMs': round(ttft * 1000) if ttft else None, 'totalMs': round(total * 1000),
            'inTok': pin, 'outTok': out,
            'tokPerSec': round(out / total, 1) if out else 0,
            'inr': round((pin * inr[0] + out * inr[1]) / 1e6, 3) if usage else 0}


ap = argparse.ArgumentParser()
ap.add_argument('--runs', type=int, default=20)
ap.add_argument('--dry', action='store_true')
ap.add_argument('--hosts', default=None)
ap.add_argument('--at', default=None)
a = ap.parse_args()
runs = 1 if a.dry else a.runs
only = a.hosts.lower().split(',') if a.hosts else None
hosts = [h for h in ALL_HOSTS if not only or h[0].lower() in only]
if not hosts:
    sys.exit('no hosts match --hosts filter')

results = {h[0]: [] for h in hosts}
for i in range(runs):
    for name, body, inr in hosts:   # round-robin
        r = one(body, inr)
        results[name].append(r)
        msg = f'ERR {r["err"]}' if r['err'] else (
            f'{r["totalMs"]}ms ttft={r["ttftMs"]}ms {r["outTok"]}tok {r["tokPerSec"]}tok/s ₹{r["inr"]}')
        print(f'[{i + 1}/{runs}] {name:<10} {msg}', file=sys.stderr, flush=True)

summary = []
for name, body, inr in hosts:
    rs = results[name]
    ok = [r for r in rs if not r['err']]
    norm = [r['totalMs'] * (700 / r['outTok']) if r['outTok'] else r['totalMs'] for r in ok]
    summary.append({
        'host': name, 'runs': len(rs), 'errors': len(rs) - len(ok),
        'medianLatencyMs': round(statistics.median([r['totalMs'] for r in ok])) if ok else 0,
        'medianLatency700Ms': round(statistics.median(norm)) if ok else 0,
        'medianTtftMs': round(statistics.median([r['ttftMs'] or 0 for r in ok])) if ok else 0,
        'medianTokPerSec': statistics.median([r['tokPerSec'] for r in ok]) if ok else 0,
        'medianOutTok': statistics.median([r['outTok'] for r in ok]) if ok else 0,
        'inrPerTask': statistics.median([r['inr'] for r in ok]) if ok else 0,
    })

payload = {'at': a.at, 'vantage': os.uname().nodename, 'runsPerHost': runs, 'prompt': PROMPT,
           'summary': summary, 'raw': results}
with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'results-glm-host-sweep-py.json'), 'w') as f:
    json.dump(payload, f, indent=1)
print(json.dumps(summary, indent=2))