from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

from check_inference_endpoint import check_endpoint


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


def main() -> int:
    parser = argparse.ArgumentParser(description="Wait for an OpenAI-compatible inference endpoint.")
    parser.add_argument("--base-url", required=True)
    parser.add_argument("--expected-model")
    parser.add_argument("--timeout", type=int, default=900)
    parser.add_argument("--interval", type=int, default=15)
    parser.add_argument("--request-timeout", type=int, default=10)
    parser.add_argument("--report", type=Path, default=DEFAULT_REPORT)
    parser.add_argument("--print", action="store_true")
    args = parser.parse_args()

    deadline = time.monotonic() + args.timeout
    attempts = []
    while True:
        result = check_endpoint(args.base_url, args.expected_model, args.request_timeout)
        result["attempt"] = len(attempts) + 1
        attempts.append(result)
        if result["status"] == "ok":
            report = {"status": "ok", "attempts": attempts}
            break
        if time.monotonic() >= deadline:
            report = {"status": "failed", "attempts": attempts}
            break
        time.sleep(args.interval)

    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 wait status: {report['status']} ({args.base_url})")
    return 0 if report["status"] == "ok" else 1


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