from __future__ import annotations

import argparse
import datetime as dt
import json
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_REPORT = ROOT / "reports" / "inference-endpoint-check.json"


def fetch_models(base_url: str, timeout: int) -> tuple[list[str], int]:
    started_at = time.perf_counter()
    request = urllib.request.Request(f"{base_url.rstrip('/')}/v1/models", method="GET")
    with urllib.request.urlopen(request, timeout=timeout) as response:
        payload = json.loads(response.read().decode("utf-8"))
    latency_ms = round((time.perf_counter() - started_at) * 1000)
    rows = payload.get("data") if isinstance(payload, dict) else []
    if not isinstance(rows, list):
        return [], latency_ms
    return [str(row.get("id")) for row in rows if isinstance(row, dict) and row.get("id")], latency_ms


def check_endpoint(base_url: str, expected_model: str | None, timeout: int) -> dict:
    result = {
        "base_url": base_url,
        "expected_model": expected_model,
        "status": "failed",
        "models": [],
        "available": None,
        "latency_ms": None,
        "error": None,
    }
    try:
        models, latency_ms = fetch_models(base_url, timeout)
    except (TimeoutError, OSError, urllib.error.URLError, ValueError) as exc:
        result["error"] = str(exc)
        return result

    result["models"] = models
    result["latency_ms"] = latency_ms
    result["available"] = expected_model in models if expected_model else None
    result["status"] = "ok" if models and (expected_model is None or result["available"]) else "failed"
    return result


def main() -> int:
    parser = argparse.ArgumentParser(description="Check an OpenAI-compatible inference endpoint.")
    parser.add_argument("--base-url", default="http://docker-gpu.cin.su:8000")
    parser.add_argument("--expected-model")
    parser.add_argument("--timeout", type=int, default=10)
    parser.add_argument("--report", type=Path, default=DEFAULT_REPORT)
    parser.add_argument("--print", action="store_true")
    parser.add_argument("--no-report", action="store_true")
    args = parser.parse_args()

    result = check_endpoint(args.base_url, args.expected_model, args.timeout)
    report = {
        "created_at": dt.datetime.now(dt.UTC).isoformat(),
        "status": result["status"],
        "checks": [result],
    }

    if not args.no_report:
        args.report.parent.mkdir(parents=True, exist_ok=True)
        args.report.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")

    if args.print:
        print(json.dumps(report, ensure_ascii=False, indent=2))
    else:
        print(f"Endpoint status: {result['status']} ({args.base_url})")
        if result["error"]:
            print(result["error"], file=sys.stderr)

    return 0 if result["status"] == "ok" else 1


if __name__ == "__main__":
    raise SystemExit(main())