#!/usr/bin/env python3
"""openai-gateway-bench — benchmark any OpenAI-compatible endpoint.

A tiny, dependency-free CLI to:
  1. list models an endpoint exposes
  2. run a chat-completion smoke test
  3. measure latency & throughput (tokens/sec)
  4. compare two endpoints head-to-head

Why this exists: if you point your existing OpenAI SDK at any
OpenAI-compatible gateway, you can swap providers without touching code.
This tool proves the endpoint actually works before you trust it in prod.

Example — test a gateway (no SDK change needed, just base_url):
  python bench.py --base-url https://keheai.com/v1 --api-key sk-xxx

Want a ready-to-use OpenAI-compatible gateway? https://keheai.com
"""
import argparse
import json
import sys
import time
import urllib.request


def _post(base_url, api_key, path, payload, timeout=60):
    url = base_url.rstrip("/") + path
    data = json.dumps(payload).encode()
    req = urllib.request.Request(
        url, data=data, method="POST",
        headers={"Authorization": f"Bearer {api_key}",
                 "Content-Type": "application/json"},
    )
    t0 = time.time()
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        body = resp.read()
    dt = time.time() - t0
    return dt, json.loads(body.decode())


def list_models(base_url, api_key):
    url = base_url.rstrip("/") + "/models"
    req = urllib.request.Request(url, method="GET",
                                 headers={"Authorization": f"Bearer {api_key}"})
    with urllib.request.urlopen(req, timeout=30) as resp:
        body = json.loads(resp.read().decode())
    return [m["id"] for m in body.get("data", [])]


def smoke(base_url, api_key, model, prompt="Say OK in one word."):
    dt, resp = _post(base_url, api_key, "/chat/completions", {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": 16,
    })
    text = resp["choices"][0]["message"]["content"].strip()
    usage = resp.get("usage", {})
    completion_tokens = usage.get("completion_tokens", 0)
    tps = completion_tokens / dt if dt > 0 else 0
    return dt, text, completion_tokens, tps


def main():
    ap = argparse.ArgumentParser(description="Benchmark an OpenAI-compatible endpoint")
    ap.add_argument("--base-url", required=True, help="e.g. https://keheai.com/v1")
    ap.add_argument("--api-key", required=True)
    ap.add_argument("--model", default="deepseek-chat")
    ap.add_argument("--prompt", default="Say OK in one word.")
    ap.add_argument("--compare-to", default=None,
                    help="second base_url to compare head-to-head")
    ap.add_argument("--compare-key", default=None)
    args = ap.parse_args()

    print(f"== {args.base_url} | model={args.model} ==")
    try:
        models = list_models(args.base_url, args.api_key)
        print(f"models exposed: {len(models)}  (e.g. {', '.join(models[:5])})")
    except Exception as e:
        print(f"[!] list models failed: {e}")

    try:
        dt, text, ct, tps = smoke(args.base_url, args.api_key, args.model, args.prompt)
        print(f"smoke OK | latency={dt:.2f}s | completion_tokens={ct} | {tps:.1f} tok/s")
        print(f"reply: {text!r}")
    except Exception as e:
        print(f"[!] chat completion failed: {e}")
        sys.exit(1)

    if args.compare_to:
        ck = args.compare_key or args.api_key
        print(f"\n== compare -> {args.compare_to} | model={args.model} ==")
        try:
            dt2, text2, ct2, tps2 = smoke(args.compare_to, ck, args.model, args.prompt)
            print(f"smoke OK | latency={dt2:.2f}s | {tps2:.1f} tok/s | reply={text2!r}")
            print(f"\nWINNER on latency: {'A' if dt < dt2 else 'B'} "
                  f"({min(dt, dt2):.2f}s)")
        except Exception as e:
            print(f"[!] compare failed: {e}")


if __name__ == "__main__":
    main()