#!/usr/bin/env python3
"""Validate a running hybrid GLM-5.2 vLLM server.

Checks: completion coherence (short prompt), decode throughput, per-GPU
memory versus the 96GB target ceiling, and reported KV-cache capacity.

Usage: validate_serve.py [--port 8199] [--long N_TOKENS]
"""
import argparse
import json
import subprocess
import time
import urllib.request


def post(port, path, body):
    req = urllib.request.Request(
        f"http://localhost:{port}{path}",
        data=json.dumps(body).encode(),
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(req, timeout=1800) as r:
        return json.load(r)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=8199)
    ap.add_argument("--long", type=int, default=0,
                    help="also run a long-context probe of this many tokens")
    args = ap.parse_args()

    models = json.load(
        urllib.request.urlopen(f"http://localhost:{args.port}/v1/models")
    )
    model = models["data"][0]["id"]
    print("model:", model)

    prompts = [
        "The capital of France is",
        "def fibonacci(n):\n    ",
        "Water is composed of the elements",
    ]
    for p in prompts:
        t0 = time.time()
        out = post(args.port, "/v1/completions", {
            "model": model, "prompt": p, "max_tokens": 48, "temperature": 0.0,
        })
        text = out["choices"][0]["text"]
        n = out["usage"]["completion_tokens"]
        dt = time.time() - t0
        print(f"\n>>> {p!r}\n{text!r}\n[{n} tok in {dt:.1f}s = {n/dt:.1f} tok/s]")

    if args.long:
        filler = ("The quick brown fox jumps over the lazy dog. " * 8 + "\n")
        target_words = int(args.long * 0.75)
        doc = (filler * (target_words // len(filler.split()) + 1))
        needle = "\nIMPORTANT: the secret code word is BLUEBERRY42.\n"
        prompt = (doc[: len(doc) // 2] + needle + doc[len(doc) // 2:]
                  + "\n\nWhat is the secret code word mentioned above? Answer:")
        t0 = time.time()
        out = post(args.port, "/v1/completions", {
            "model": model, "prompt": prompt, "max_tokens": 16,
            "temperature": 0.0,
        })
        dt = time.time() - t0
        u = out["usage"]
        print(f"\n=== long-context probe: {u['prompt_tokens']} prompt tokens, "
              f"{dt:.0f}s ===")
        print("answer:", out["choices"][0]["text"].strip())

    print("\n=== per-GPU memory ===")
    smi = subprocess.run(
        ["nvidia-smi", "--query-gpu=index,memory.used",
         "--format=csv,noheader,nounits"],
        capture_output=True, text=True,
    ).stdout
    worst = 0
    for line in smi.strip().splitlines():
        idx, used = [int(x) for x in line.split(",")]
        worst = max(worst, used)
        flag = "OK" if used <= 97887 else "OVER 96GB TARGET"
        print(f"gpu{idx}: {used} MiB  [{flag}]")
    print(f"worst GPU: {worst} MiB vs 97887 MiB target "
          f"({'FITS' if worst <= 97887 else 'DOES NOT FIT'})")


if __name__ == "__main__":
    main()