#!/usr/bin/env python3
#
#  verify.py
#  Zyquo Router
#
#  Author: Simon-Pierre Boucher
#  Mail: contact@spboucher.ai
#
#  Phase 7 verification harness. Drives EVERY chat model in the catalog
#  through the router's local endpoint using the official OpenAI Python SDK
#  (never the upstreams directly): non-streaming, streaming (SDK-parsed
#  chunk discipline), tools / vision / reasoning where the catalog says the
#  model supports them. Emits the compatibility matrix as Markdown.
#
#  Usage:  verify.py [--base http://127.0.0.1:8787/v1] [--providers xai,mistral]
#          [--models id1,id2] [--workers 6] [--out docs/VERIFICATION.md]
#
import argparse
import base64
import concurrent.futures
import json
import sys
import threading
import time
import urllib.request

from openai import OpenAI, APIError, APIStatusError

TOOLS = [{
    "type": "function",
    "function": {
        "name": "get_weather",
        "description": "Get current weather for a location",
        "parameters": {
            "type": "object",
            "properties": {"location": {"type": "string"}},
            "required": ["location"],
        },
    },
}]

# 64×64 solid red PNG (1×1 images are rejected by some providers).
TINY_PNG = "iVBORw0KGgoAAAANSUhEUgAAAEAAAABACAIAAAAlC+aJAAAAb0lEQVR4nO3PAQkAAAyEwO9feoshgnABdLep8QUNyPEFDcjxBQ3I8QUNyPEFDcjxBQ3I8QUNyPEFDcjxBQ3I8QUNyPEFDcjxBQ3I8QUNyPEFDcjxBQ3I8QUNyPEFDcjxBQ3I8QUNyPEFDcjxBQ3IPanc8OLDQitxAAAAAElFTkSuQmCC"

print_lock = threading.Lock()


def log(msg):
    with print_lock:
        print(msg, flush=True)


class ModelResult:
    def __init__(self, model_id, meta):
        self.id = model_id
        self.meta = meta or {}
        self.non_stream = None   # "PASS" / "FAIL: …" / "RATE-LIMITED"
        self.stream = None
        self.tools = "n/a"
        self.vision = "n/a"
        self.reasoning = "n/a"

    @property
    def ok(self):
        checks = [self.non_stream, self.stream, self.tools, self.vision, self.reasoning]
        return all(c in (None, "n/a") or str(c).startswith(("PASS", "RATE-LIMITED", "SKIP")) for c in checks)


def classify(err):
    if isinstance(err, APIStatusError) and err.status_code == 429:
        return "RATE-LIMITED"
    text = str(err)
    return "FAIL: " + text[:160].replace("\n", " ").replace("|", "/")


def check_non_stream(client, result):
    r = client.chat.completions.create(
        model=result.id, max_tokens=32,
        messages=[{"role": "user", "content": "Reply with exactly: OK"}],
    )
    assert r.object == "chat.completion", f"object={r.object}"
    assert r.model == result.id, f"model echo {r.model}"
    assert r.choices[0].message.role == "assistant"
    assert r.choices[0].finish_reason in ("stop", "length"), f"finish={r.choices[0].finish_reason}"
    assert r.usage and r.usage.total_tokens > 0, "usage missing"
    return "PASS"


def check_stream(client, result):
    text, finish, usage, got_role = "", None, None, False
    stream = client.chat.completions.create(
        model=result.id, max_tokens=64, stream=True,
        stream_options={"include_usage": True},
        messages=[{"role": "user", "content": "Count from 1 to 3, digits only."}],
    )
    first = True
    for chunk in stream:
        assert chunk.object == "chat.completion.chunk", f"chunk object={chunk.object}"
        if chunk.usage:
            usage = chunk.usage
        if not chunk.choices:
            continue
        delta = chunk.choices[0].delta
        if first and delta.role == "assistant":
            got_role = True
        first = False
        if chunk.choices[0].finish_reason:
            finish = chunk.choices[0].finish_reason
        text += delta.content or ""
    assert got_role, "no role delta on first chunk"
    assert finish in ("stop", "length"), f"finish={finish}"
    assert usage is not None and usage.total_tokens > 0, "usage chunk missing"
    return "PASS"


def check_tools(client, result, tool_choice="auto"):
    calls = {}
    stream = client.chat.completions.create(
        model=result.id, max_tokens=300, stream=True,
        tools=TOOLS, tool_choice=tool_choice,
        messages=[{"role": "user", "content": "What's the weather in Paris? Use the get_weather tool."}],
    )
    for chunk in stream:
        if not chunk.choices:
            continue
        for tc in chunk.choices[0].delta.tool_calls or []:
            entry = calls.setdefault(tc.index, {"id": None, "name": "", "args": ""})
            if tc.id:
                entry["id"] = tc.id
            if tc.function and tc.function.name:
                entry["name"] = tc.function.name
            if tc.function and tc.function.arguments:
                entry["args"] += tc.function.arguments
    assert calls, "no tool call streamed"
    call = calls[min(calls)]
    assert call["id"], "tool call id missing"
    assert call["name"] == "get_weather", f"name={call['name']}"
    try:
        args = json.loads(call["args"])
        note = ""
    except json.JSONDecodeError:
        # DashScope qwq repeats the complete arguments object per delta —
        # accept the first object but flag the upstream quirk.
        args, _ = json.JSONDecoder().raw_decode(call["args"])
        note = " (upstream repeats args)"
    assert "location" in args, f"args={args}"
    return "PASS" + note


def check_vision(client, result):
    r = client.chat.completions.create(
        model=result.id, max_tokens=1024,
        messages=[{"role": "user", "content": [
            {"type": "text", "text": "One word: what color is this image?"},
            {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{TINY_PNG}"}},
        ]}],
    )
    content = r.choices[0].message.content or ""
    assert content.strip(), "empty vision answer"
    return "PASS"


def check_reasoning(client, result):
    r = client.chat.completions.create(
        model=result.id, max_tokens=2048,
        extra_body={"reasoning_effort": "low"},
        messages=[{"role": "user", "content": "What is 17*23? Reply with the number only."}],
    )
    message = r.choices[0].message
    reasoning = getattr(message, "reasoning_content", None)
    if reasoning is None and message.model_extra:
        reasoning = message.model_extra.get("reasoning_content")
    details = (r.usage.completion_tokens_details if r.usage else None)
    reasoning_tokens = getattr(details, "reasoning_tokens", None) if details else None
    if (reasoning and reasoning.strip()) or (reasoning_tokens or 0) > 0:
        return "PASS"
    # The call accepted reasoning_effort and answered, but the provider keeps
    # reasoning server-side (OpenAI o-series/gpt-5*, some hosted models).
    assert (r.choices[0].message.content or "").strip(), "empty reasoning answer"
    return "PASS (hidden)"


def verify_model(base, entry):
    client = OpenAI(base_url=base, api_key="zyquo-verify", timeout=180, max_retries=0)
    meta = entry.get("x_zyquo") or {}
    result = ModelResult(entry["id"], meta)

    # Deep-research models run for minutes — beyond any sane harness timeout.
    long_running = "deep-research" in result.id

    for name, fn, gated in [
        ("non_stream", check_non_stream, True),
        ("stream", check_stream, not long_running),
        ("tools", check_tools, meta.get("tools")),
        ("vision", check_vision, meta.get("vision")),
        ("reasoning", check_reasoning, meta.get("reasoning") and not long_running),
    ]:
        if not gated:
            if long_running and name in ("stream", "reasoning"):
                setattr(result, name, "SKIP (long-running)")
            continue
        try:
            setattr(result, name, fn(client, result))
        except AssertionError as err:
            outcome = f"FAIL: {err}"
            # Weak tool-callers may ignore "auto" — one retry forcing the call.
            if name == "tools" and "no tool call" in str(err):
                try:
                    outcome = check_tools(client, result, tool_choice="required") + " (required)"
                except Exception as retry_err:  # noqa: BLE001
                    outcome = f"FAIL: no tool call with auto; required → {classify(retry_err)[:80]}"
            setattr(result, name, outcome)
        except (APIError, Exception) as err:  # noqa: BLE001 — harness must not die
            setattr(result, name, classify(err))
        time.sleep(0.3)

    status = "OK " if result.ok else "!! "
    log(f"{status}{result.id}: ns={result.non_stream} st={result.stream} "
        f"tools={result.tools} vision={result.vision} reasoning={result.reasoning}")
    return result


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--base", default="http://127.0.0.1:8787/v1")
    parser.add_argument("--providers", default="")
    parser.add_argument("--models", default="")
    parser.add_argument("--workers", type=int, default=6)
    parser.add_argument("--out", default="docs/VERIFICATION.md")
    args = parser.parse_args()

    with urllib.request.urlopen(f"{args.base}/models") as response:
        catalog = json.load(response)["data"]

    if args.providers:
        wanted = set(args.providers.split(","))
        catalog = [m for m in catalog if m["owned_by"] in wanted]
    if args.models:
        wanted = set(args.models.split(","))
        catalog = [m for m in catalog if m["id"] in wanted]

    log(f"Verifying {len(catalog)} models via {args.base}")
    results = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as pool:
        futures = [pool.submit(verify_model, args.base, entry) for entry in catalog]
        for future in concurrent.futures.as_completed(futures):
            results.append(future.result())

    results.sort(key=lambda r: r.id)
    failed = [r for r in results if not r.ok]

    lines = [
        "# Zyquo Router — Phase 7 Compatibility Matrix",
        "",
        f"Generated {time.strftime('%Y-%m-%d %H:%M')} by `scripts/verify.py` — every request",
        "went through the router's local endpoint using the official OpenAI Python SDK.",
        "",
        f"**{len(results)} models · {len(results) - len(failed)} green · {len(failed)} failing**",
        "",
        "| Model | Non-stream | Stream | Tools | Vision | Reasoning |",
        "|---|---|---|---|---|---|",
    ]
    for r in results:
        lines.append(f"| `{r.id}` | {r.non_stream} | {r.stream} | {r.tools} | {r.vision} | {r.reasoning} |")
    with open(args.out, "w") as handle:
        handle.write("\n".join(lines) + "\n")
    log(f"\n{len(results) - len(failed)}/{len(results)} green → {args.out}")
    sys.exit(1 if failed else 0)


if __name__ == "__main__":
    main()