"""Benchmark real do endpoint: TTFT, prompt/s, gen/s.

Mede com o objeto `timings` que o llama-server devolve (nao estima por
relogio de parede) e registra tudo em JSON para comparar quantizacoes e
contagens de threads.

Uso:
  BASE=https://erinaldordso-programa.hf.space/v1 python3 tests/benchmark.py
  BASE=... MODEL=... APIKEY=... CTX=32768 python3 tests/benchmark.py

Prompts: um curto (TTFT puro) e um longo (~3k tokens, prefill real).
"""
import json
import os
import time
import urllib.request

BASE = os.environ.get("BASE", "https://erinaldordso-programa.hf.space/v1").rstrip("/")
MODEL = os.environ.get("MODEL", "qwen3-coder-30b-a3b-instruct")
APIKEY = os.environ.get("APIKEY", "")
CTX = os.environ.get("CTX", "32768")

LONG_PROMPT = ("Explique o padrão de projeto observador em Python. " * 120
               + "\nMostre uma implementação completa com testes.")


def post(body):
    headers = {"Content-Type": "application/json"}
    if APIKEY:
        headers["Authorization"] = f"Bearer {APIKEY}"
    t0 = time.perf_counter()
    r = urllib.request.Request(BASE + "/chat/completions",
                               data=json.dumps(body).encode(),
                               headers=headers, method="POST")
    first = None
    chunks = 0
    with urllib.request.urlopen(r, timeout=600) as resp:
        for line in resp:
            s = line.decode(errors="replace").strip()
            if not s.startswith("data:"):
                continue
            s = s[5:].strip()
            if s == "[DONE]":
                last = {}
                break
            try:
                last = json.loads(s)
            except json.JSONDecodeError:
                continue
            chunks += 1
            if first is None:
                first = time.perf_counter() - t0
    total = time.perf_counter() - t0
    return last, {"ttft_s": round(first or 0, 3), "total_s": round(total, 2),
                  "chunks": chunks}


def run_case(name, prompt, max_tokens=128):
    last, wall = post({"model": MODEL,
                       "messages": [{"role": "user", "content": prompt}],
                       "max_tokens": max_tokens, "temperature": 0, "stream": True,
                       "timings_per_token": True})
    usage = last.get("usage", {})
    timings = last.get("timings", {})
    row = {"case": name, "model": MODEL, "ctx": CTX,
           "prompt_tokens": usage.get("prompt_tokens"),
           "completion_tokens": usage.get("completion_tokens"),
           "prompt_per_second": timings.get("prompt_per_second"),
           "predicted_per_second": timings.get("predicted_per_second"),
           **wall}
    print(json.dumps(row, indent=2))
    return row


def main():
    print(f"BASE={BASE} MODEL={MODEL}")
    rows = [run_case("curto", "diga ok", 32),
            run_case("longo-3k", LONG_PROMPT, 128)]
    fname = f"benchmark_{MODEL}_{int(time.time())}.json"
    with open(fname, "w") as fh:
        json.dump(rows, fh, indent=2)
    print(f"salvo em {fname}")


if __name__ == "__main__":
    main()