"""Check one running `vllm-sr-runtime` instance against its pins and reference answers.

Runs inside the runtime image on the private runtime network. It checks the
served artifact (repository, revision, manifest and model identity, exact
profile, readiness), asks the runtime's own golden request twice on the exact
profile and compares the answers with the recorded reference answers for the
device class: `vllm_sr_runtime/registry/golden_answers.json` of a
vllm-project/semantic-router commit that records the same revision.

    python3 check_runtime.py --origin http://172.31.231.11:8100 \
      --model vllm-sr/Decision-2.0-Kai-0.6B --revision <40-hex> \
      --manifest-sha256 <64-hex> --model-sha256 <64-hex> --reference golden_answers.json
"""

from __future__ import annotations

import argparse
import json
import sys
import time
import urllib.request

from vllm_sr_runtime.families.decision2.family import GOLDEN_QUESTIONS, GOLDEN_STATE


def call(origin: str, path: str, body: dict | None = None) -> tuple[int, dict, float]:
    data = None if body is None else json.dumps(body).encode()
    request = urllib.request.Request(
        origin + path, data=data, method="GET" if body is None else "POST",
        headers={"Content-Type": "application/json"},
    )
    started = time.perf_counter()
    try:
        with urllib.request.urlopen(request, timeout=60) as response:
            payload = json.load(response)
            status = response.status
    except urllib.error.HTTPError as error:
        payload, status = json.load(error), error.code
    return status, payload, (time.perf_counter() - started) * 1000


def max_delta(left: dict, right: dict) -> float:
    worst = 0.0
    for key, answer in right.items():
        mine = left.get(key, {})
        if answer.get("type") == "noul":
            worst = max(worst, abs(mine.get("noul", 2.0) - answer["noul"]))
            continue
        for candidate, value in answer.get("probabilities", {}).items():
            worst = max(worst, abs(mine.get("probabilities", {}).get(candidate, 2.0) - value))
    return worst


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--origin", required=True)
    parser.add_argument("--model", required=True)
    parser.add_argument("--revision", required=True)
    parser.add_argument("--manifest-sha256", required=True)
    parser.add_argument("--model-sha256", required=True)
    parser.add_argument("--reference", required=True)
    parser.add_argument("--device-class", default="rocm")
    args = parser.parse_args()

    _, health, _ = call(args.origin, "/health")
    _, listing, _ = call(args.origin, "/v1/models")
    card = listing["data"][0]
    identity = {
        "ready": health.get("status") == "ready" and card.get("ready") is True,
        "repo": card.get("repo") == args.model and card.get("id") == args.model,
        "revision": card.get("revision") == args.revision,
        "manifest_sha256": card.get("manifest_sha256") == args.manifest_sha256,
        "model_sha256": card.get("model_sha256") == args.model_sha256,
        "profile": card.get("profile") == "exact",
    }
    body = {"model": args.model, "state": GOLDEN_STATE, "questions": GOLDEN_QUESTIONS,
            "options": {"profile": "exact", "return_meta": True}}
    status_a, first, ms_a = call(args.origin, "/v1/systemone", body)
    status_b, second, ms_b = call(args.origin, "/v1/systemone", body)
    recorded = json.load(open(args.reference, encoding="utf-8")).get(args.model, {})
    reference = (recorded.get("answers") or {}).get(args.device_class)
    answers = first.get("answers", {})
    report = {
        "model": args.model,
        "identity": identity,
        "device": card.get("device"),
        "accelerator_validated": card.get("accelerator_validated"),
        "golden_status": (card.get("golden") or {}).get("status"),
        "deterministic": status_a == status_b == 200 and first.get("answers") == second.get("answers"),
        "meta_revision": first.get("meta", {}).get("revision") == args.revision,
        "reference_revision": recorded.get("revision"),
        "reference_identical": reference is not None and answers == reference,
        "reference_max_abs_delta": None if reference is None else max_delta(answers, reference),
        "latency_ms": [round(ms_a, 2), round(ms_b, 2)],
    }
    print(json.dumps(report, sort_keys=True))
    ok = (all(identity.values()) and report["deterministic"] and report["meta_revision"]
          and reference is not None and report["reference_max_abs_delta"] <= 0.02)
    return 0 if ok else 1


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