#!/usr/bin/env python3
"""Pre-cutover surface test: every capability the live gen seat actually serves.

The gen seat backs 7 LiteLLM aliases (gen, gen-reasoning, summarizer,
summarizer-large, classifier, image-judge, qwen-image-bench), so a cutover has
to clear vision, tool-calling, the thinking split, long context, and streaming --
not just decode speed.
"""
import base64, json, struct, sys, urllib.request, zlib, argparse

def rpc(base, path, payload, timeout=900):
    req = urllib.request.Request(base + path, data=json.dumps(payload).encode(),
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.load(r)

def png(w, h, fn):
    raw = b"".join(b"\x00" + bytes(v for x in range(w) for v in fn(x, y)) for y in range(h))
    def chunk(t, d):
        c = t + d
        return struct.pack(">I", len(d)) + c + struct.pack(">I", zlib.crc32(c) & 0xffffffff)
    return (b"\x89PNG\r\n\x1a\n"
            + chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0))
            + chunk(b"IDAT", zlib.compress(raw)) + chunk(b"IEND", b""))

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", required=True)
    ap.add_argument("--model", required=True)
    ap.add_argument("--thinking-model", default=None)
    a = ap.parse_args()
    B, M = a.base, a.model
    results = []

    def check(name, ok, detail=""):
        results.append((name, ok, detail))
        print(f"  [{'PASS' if ok else 'FAIL'}] {name}: {detail[:150]}")

    # 1. plain chat
    try:
        r = rpc(B, "/v1/chat/completions", {"model": M, "max_tokens": 60, "temperature": 0,
              "messages": [{"role": "user", "content": "Name the largest moon of Saturn in one word."}],
              "chat_template_kwargs": {"enable_thinking": False}})
        c = (r["choices"][0]["message"].get("content") or "")
        check("plain chat", "titan" in c.lower(), repr(c.strip()))
    except Exception as e:
        check("plain chat", False, str(e))

    # 2. vision
    try:
        img = png(64, 64, lambda x, y: (30, 90, 220) if (14 <= x < 50 and 14 <= y < 50) else (250, 250, 250))
        b64 = base64.b64encode(img).decode()
        r = rpc(B, "/v1/chat/completions", {"model": M, "max_tokens": 60, "temperature": 0,
              "messages": [{"role": "user", "content": [
                  {"type": "text", "text": "What colour is the square in this image? One word."},
                  {"type": "image_url", "image_url": {"url": "data:image/png;base64," + b64}}]}],
              "chat_template_kwargs": {"enable_thinking": False}})
        c = (r["choices"][0]["message"].get("content") or "")
        check("vision (image)", "blue" in c.lower(), repr(c.strip()))
    except Exception as e:
        check("vision (image)", False, str(e))

    # 3. tool calling
    try:
        r = rpc(B, "/v1/chat/completions", {"model": M, "max_tokens": 200, "temperature": 0,
              "messages": [{"role": "user", "content": "What's the weather in Anaheim? Use the tool."}],
              "tools": [{"type": "function", "function": {"name": "get_weather",
                  "description": "Get current weather for a city",
                  "parameters": {"type": "object", "properties": {"city": {"type": "string"}},
                                 "required": ["city"]}}}]})
        tc = r["choices"][0]["message"].get("tool_calls")
        ok = bool(tc) and tc[0]["function"]["name"] == "get_weather" and "Anaheim" in tc[0]["function"]["arguments"]
        check("tool calling", ok, json.dumps(tc)[:150] if tc else "no tool_calls")
    except Exception as e:
        check("tool calling", False, str(e))

    # 4. thinking split (reasoning parser)
    tm = a.thinking_model or M
    try:
        r = rpc(B, "/v1/chat/completions", {"model": tm, "max_tokens": 400, "temperature": 0,
              "messages": [{"role": "user", "content": "A bat and ball cost $1.10 total. The bat costs $1 more than the ball. What does the ball cost?"}],
              "chat_template_kwargs": {"enable_thinking": True}})
        m = r["choices"][0]["message"]
        rc = m.get("reasoning") or m.get("reasoning_content") or ""
        c = m.get("content") or ""
        check("thinking split", len(rc) > 0 or len(c) > 0,
              f"reasoning={len(rc)}ch content={len(c)}ch :: {(c or rc)[:80]!r}")
    except Exception as e:
        check("thinking split", False, str(e))

    # 5. long context (~40k tokens, well past the 32k probe ceiling)
    try:
        filler = "The archived maintenance log records routine inspection of pump assembly seven. " * 3000
        needle = "\n\nIMPORTANT: the calibration passphrase is HELIOTROPE-49.\n\n"
        prompt = filler[:len(filler)//2] + needle + filler[len(filler)//2:] + \
                 "\n\nWhat is the calibration passphrase? Answer with just the passphrase."
        r = rpc(B, "/v1/chat/completions", {"model": M, "max_tokens": 40, "temperature": 0,
              "messages": [{"role": "user", "content": prompt}],
              "chat_template_kwargs": {"enable_thinking": False}})
        c = (r["choices"][0]["message"].get("content") or "")
        pt = r["usage"]["prompt_tokens"]
        check("long context + retrieval", "HELIOTROPE-49" in c.upper(),
              f"{pt} prompt tokens -> {c.strip()[:60]!r}")
    except Exception as e:
        check("long context + retrieval", False, str(e))

    # 6. streaming
    try:
        req = urllib.request.Request(B + "/v1/chat/completions",
            data=json.dumps({"model": M, "max_tokens": 60, "temperature": 0, "stream": True,
                "messages": [{"role": "user", "content": "Count from 1 to 5."}],
                "chat_template_kwargs": {"enable_thinking": False}}).encode(),
            headers={"Content-Type": "application/json"})
        n = 0
        with urllib.request.urlopen(req, timeout=300) as resp:
            for line in resp:
                if line.startswith(b"data: ") and b"[DONE]" not in line:
                    n += 1
        check("streaming", n > 3, f"{n} SSE chunks")
    except Exception as e:
        check("streaming", False, str(e))

    npass = sum(1 for _, ok, _ in results if ok)
    print(f"\n{npass}/{len(results)} passed")
    return 0 if npass == len(results) else 1

if __name__ == "__main__":
    sys.exit(main())