#!/usr/bin/env python3
"""尝试通过 NVIDIA Sync 连接 DGXA，直到 LLM3D capabilities 绿灯。"""
from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
import time
from pathlib import Path
from urllib.request import urlopen

DEFAULT_BASE = os.environ.get("LLM3D_BASE", "http://192.168.130.16:8008").rstrip("/")
NVSYNC = Path(
    os.environ.get(
        "NVSYNC_EXE",
        str(
            Path(os.environ.get("LOCALAPPDATA", ""))
            / "Programs"
            / "nvidia-sync"
            / "resources"
            / "bin"
            / "nvsync-amd64.exe"
        ),
    )
)


def caps(base: str) -> dict:
    with urlopen(f"{base}/api/capabilities", timeout=90) as r:
        return json.loads(r.read().decode("utf-8"))


def is_ready(c: dict) -> bool:
    dgx = c.get("dgx") or {}
    return bool(dgx.get("tokenrig_ready") or dgx.get("unirig_ready")) and bool(
        dgx.get("unimate_ready")
    )


def try_nvsync_connect(alias: str = "DGXA") -> None:
    if not NVSYNC.is_file():
        print(f"nvsync missing: {NVSYNC}", flush=True)
        return
    print(f"nvsync connect {alias} --detach ...", flush=True)
    try:
        subprocess.run(
            [str(NVSYNC), "connect", alias, "--detach"],
            timeout=90,
            check=False,
            capture_output=True,
            text=True,
        )
    except subprocess.TimeoutExpired:
        print("nvsync connect timed out (device may be offline)", flush=True)
    st = subprocess.run(
        [str(NVSYNC), "status", alias],
        timeout=30,
        check=False,
        capture_output=True,
        text=True,
    )
    print(st.stdout or st.stderr or "(no status)", flush=True)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", default=DEFAULT_BASE)
    ap.add_argument("--alias", default="DGXA")
    ap.add_argument("--retries", type=int, default=30)
    ap.add_argument("--interval", type=float, default=20.0)
    ap.add_argument("--no-connect", action="store_true")
    args = ap.parse_args()

    for i in range(1, args.retries + 1):
        print(f"=== attempt {i}/{args.retries} ===", flush=True)
        if not args.no_connect:
            try_nvsync_connect(args.alias)
        try:
            c = caps(args.base)
        except Exception as exc:  # noqa: BLE001
            print(f"capabilities error: {exc}", flush=True)
            time.sleep(args.interval)
            continue
        dgx = c.get("dgx") or {}
        print(
            json.dumps(
                {
                    "tokenrig_ready": dgx.get("tokenrig_ready"),
                    "unimate_ready": dgx.get("unimate_ready"),
                    "hy3d_ready": dgx.get("hy3d_ready"),
                    "tokenrig_detail": dgx.get("tokenrig_detail"),
                },
                ensure_ascii=False,
            ),
            flush=True,
        )
        if is_ready(c):
            print("READY", flush=True)
            return 0
        time.sleep(args.interval)

    print("NOT_READY", flush=True)
    return 2


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