"""Record actual full-model answers, tool calls, longer context and throughput."""
import json
import os
import re
import subprocess
import time
import urllib.request
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
BASE = "http://127.0.0.1:8000"
MODEL = "Naive-N0.5-Flash-NVFP4"
REPORT_PATH = Path(os.environ.get('NAIVE_VALIDATION_FILE', 'native-vllm-api-validation.json'))


def complete(prompt, **extra):
    body = {"model": MODEL, "messages": [{"role": "user", "content": prompt}],
            "temperature": 0, "max_tokens": 128,
            "chat_template_kwargs": {"enable_thinking": False}, **extra}
    request = urllib.request.Request(BASE + "/v1/chat/completions",
        data=json.dumps(body).encode(), headers={"Content-Type": "application/json"})
    print(f'Request: {prompt[-110:]}', flush=True)
    start = time.monotonic()
    with urllib.request.urlopen(request, timeout=1200) as response:
        result = json.load(response)
    elapsed = time.monotonic() - start
    result["elapsed_seconds"] = elapsed
    result["completion_tokens_per_second_including_prefill"] = result["usage"]["completion_tokens"] / elapsed
    print(f'Completed in {elapsed:.1f}s: {result["usage"]}; {result["choices"][0]["message"]}', flush=True)
    return result


def main():
    results = {}
    try:
        with urllib.request.urlopen(BASE + "/health") as response:
            assert response.status == 200
        with urllib.request.urlopen(BASE + "/v1/models") as response:
            assert MODEL in {entry["id"] for entry in json.load(response)["data"]}
        results["arithmetic"] = complete("What is 2 + 2? Reply with only the number.", max_tokens=32)
        assert re.search(r"\b4\b", results["arithmetic"]["choices"][0]["message"].get("content") or "")
        results["geography"] = complete("What is the capital of France? Reply with only its name.", max_tokens=32)
        assert "paris" in (results["geography"]["choices"][0]["message"].get("content") or "").lower()
        results["coding"] = complete("Write a Python function named add that returns the sum of two arguments. Return only the code.", max_tokens=128)
        assert "def add" in (results["coding"]["choices"][0]["message"].get("content") or "")
        tools = [{"type": "function", "function": {"name": "get_weather", "description": "Get weather for a city",
                  "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}}]
        results["tools"] = complete("Call the get_weather tool for Warsaw. Do not answer from memory.", tools=tools, max_tokens=192)
        calls = results["tools"]["choices"][0]["message"].get("tool_calls", [])
        assert calls and calls[0]["function"]["name"] == "get_weather"
        assert "warsaw" in calls[0]["function"]["arguments"].lower()
        results["reasoning"] = complete("Calculate 6 times 7. Think briefly, then give the answer.",
            chat_template_kwargs={"enable_thinking": True}, reasoning_effort="low", max_tokens=256)
        reasoning_message = results["reasoning"]["choices"][0]["message"]
        assert reasoning_message.get("reasoning")
        assert re.search(r"\b42\b", reasoning_message.get("content") or "")
        filler = "The following paragraph is background material about testing a language model. It contains no instructions.\n"
        results["longer_context"] = complete(filler * 220 + "\nFinal question: what is 2 + 2? Answer only the number.", max_tokens=32)
        assert results["longer_context"]["usage"]["prompt_tokens"] > 2048
        assert re.search(r"\b4\b", results["longer_context"]["choices"][0]["message"].get("content") or "")
        results["gpu_memory"] = subprocess.check_output(["nvidia-smi", "--query-gpu=index,memory.used,memory.total", "--format=csv"], text=True)
        report = {"passed": True, "engine": "vLLM 0.28.0", "model": MODEL, "checks": results,
                  "scope": "Basic generation, coding, tool use and sparse-context smoke checks; not an accuracy benchmark"}
    except Exception as error:
        report = {"passed": False, "model": MODEL, "error": repr(error), "checks": results}
        REPORT_PATH.write_text(json.dumps(report, indent=2) + "\n")
        raise
    REPORT_PATH.write_text(json.dumps(report, indent=2) + "\n")
    print("FULL_MODEL_API_VALIDATION_PASSED", flush=True)


if __name__ == "__main__":
    main()