"""Wait until SwarmUI HTTP is up and backend is Idle (on the VM)."""

from __future__ import annotations

import time
from collections.abc import Callable

from gpu_rent.config import Config
from gpu_rent.errors import CloudError
from gpu_rent.ssh_ops import run_ssh

Log = Callable[[str], None]

_REMOTE_POLL = r"""
import json, time, urllib.request
url = "http://127.0.0.1:7801"
deadline = time.time() + 25
while time.time() < deadline:
    try:
        req = urllib.request.Request(
            url + "/API/GetNewSession",
            data=b"{}",
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=5) as resp:
            session = json.loads(resp.read().decode())
        sid = session.get("session_id")
        if not sid:
            time.sleep(2)
            continue
        body = json.dumps({"session_id": sid}).encode()
        req2 = urllib.request.Request(
            url + "/API/GetCurrentStatus",
            data=body,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(req2, timeout=5) as resp:
            data = json.loads(resp.read().decode())
        st = data.get("status") or {}
        be = data.get("backend_status") or {}
        waiting = int(st.get("waiting_gens") or 0)
        live = int(st.get("live_gens") or 0)
        loading = int(st.get("loading_models") or 0)
        bstat = str(be.get("status") or "unknown").lower()
        if waiting or live or loading:
            print(f"BUSY queue w={waiting} live={live} load={loading}")
        elif bstat not in ("idle", "disabled", "all_disabled", "empty"):
            print(f"BUSY backend={bstat}")
        else:
            print(f"READY backend={bstat}")
        raise SystemExit(0)
    except Exception as exc:
        print(f"WAIT {exc}")
        time.sleep(3)
print("WAIT timeout-slice")
"""


def wait_backend_idle(
    cfg: Config,
    host: str,
    log: Log,
    *,
    timeout: float = 2400.0,
    poll_every: float = 15.0,
) -> None:
    """Block until SwarmUI on the VM reports Idle backend (or timeout)."""
    deadline = time.time() + timeout
    log("жду HTTP :7801 и Idle backend на VM…")
    last = ""
    while time.time() < deadline:
        try:
            out = run_ssh(
                cfg,
                host,
                "python3 - <<'PY'\n" + _REMOTE_POLL + "\nPY",
                check=False,
                timeout=40,
            ).strip()
        except Exception as exc:
            out = f"WAIT ssh: {exc}"
        line = out.splitlines()[-1] if out else "WAIT empty"
        if line != last:
            log(line)
            last = line
        if line.startswith("READY"):
            log("backend Idle")
            return
        time.sleep(poll_every)
    raise CloudError(
        f"backend не стал Idle за {int(timeout)} с. "
        "Проверь journalctl -u swarmui на VM; GPU всё ещё жив."
    )