"""CPU-only live HTTP checks; generate reference once, then check each server."""
import argparse
import importlib.metadata
import json
import math
import os
from pathlib import Path
import time
import urllib.error
import urllib.request

os.environ["CUDA_VISIBLE_DEVICES"] = ""
os.environ["SAIF_DEVICE"] = "cpu"

from serve import DecisionRequest, app, load_model

ROOT = Path(__file__).resolve().parent


def requests():
    example = json.loads((ROOT / "request.json").read_text())
    return [dict(example, questions={key: value}) for key, value in example["questions"].items()]


def post(url, payload):
    request = urllib.request.Request(url + "/v1/systemone", data=json.dumps(payload).encode(),
                                     headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(request, timeout=600) as response:
        return json.load(response)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("mode", choices=["reference", "local", "container"])
    parser.add_argument("--url", default="http://127.0.0.1:18000")
    args = parser.parse_args()
    target = ROOT / "TESTS.json"
    report = json.loads(target.read_text()) if target.exists() else {}
    if args.mode == "reference":
        model = load_model()
        outputs = [model.decide(item["state"], item["questions"]) for item in requests()]
        report["reference"] = json.loads(json.dumps(outputs))
        report["versions"] = {name: importlib.metadata.version(name) for name in
                              ["torch", "transformers", "safetensors", "fastapi", "uvicorn",
                               "pydantic", "huggingface_hub", "accelerate"]}
        report["gpu_tested"] = False
        report["gpu_note"] = "CPU only; local Docker has runc only, no NVIDIA runtime."
        report["routes"] = sorted(route.path for route in app.routes)
        assert set(report["routes"]) == {"/health", "/v1/systemone", "/docs", "/docs/oauth2-redirect", "/redoc", "/openapi.json"}
        report["no_generation_routes"] = True
        DecisionRequest.model_validate(dict(requests()[0], questions={str(index): {"type": "noul"} for index in range(64)}))
        report["64_question_schema_accepted"] = True
    else:
        with urllib.request.urlopen(args.url + "/health") as response:
            assert json.load(response) == {"status": "ok", "model": "saif-4b"}
        results = []
        for payload, expected in zip(requests(), report["reference"]):
            start = time.perf_counter()
            actual = post(args.url, payload)
            elapsed = time.perf_counter() - start
            assert actual == expected, (actual, expected)
            for answer in actual["answers"].values():
                probabilities = list(answer["probabilities"].values())
                assert all(math.isfinite(value) and 0 <= value <= 1 for value in probabilities)
                assert math.isclose(sum(probabilities), 1, abs_tol=1e-12)
            results.append({"type": next(iter(payload["questions"].values()))["type"],
                            "exact_match": True, "probabilities_sum_one": True, "latency_seconds": elapsed})
        malformed = [{}, {"state": "x", "questions": {}},
                     {"state": "x", "questions": {"q": {"type": "text"}}},
                     {"state": "x", "questions": {"q": {"type": "choice", "criteria": []}}},
                     {"state": "x", "questions": {"q": {"type": "score", "criteria": {"a": "x"}}}},
                     {"state": "x", "questions": {"q": {"type": "noul", "criteria": {"bad": "x"}}}},
                     dict(requests()[0], images=["unsupported"]),
                     dict(requests()[0], questions={str(index): {"type": "noul"} for index in range(65)}),
                     dict(requests()[0], state="word " * 132000)]
        for payload in malformed:
            try:
                post(args.url, payload)
            except urllib.error.HTTPError as error:
                assert error.code == 422, error.code
            else:
                raise AssertionError("Malformed input was accepted")
        for route in ("/v1/chat/completions", "/v1/completions", "/generate"):
            try:
                urllib.request.urlopen(urllib.request.Request(args.url + route, data=b"{}"))
            except urllib.error.HTTPError as error:
                assert error.code == 404
            else:
                raise AssertionError("Unexpected generation endpoint")
        report[args.mode] = {"passed": True, "requests": results, "malformed_422_count": len(malformed),
                             "generation_routes_404": True}
    target.write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps(report.get(args.mode, {"reference": "passed"}), indent=2))


if __name__ == "__main__":
    main()