#!/usr/bin/env python3
"""Headless SwarmUI first-install: ComfyUI backend via InstallConfirmWS.

Runs on the VM (stdlib only). Idempotent: skips when backends exist or
dlbackend/ComfyUI venv is already present and IsInstalled.
"""

from __future__ import annotations

import base64
import hashlib
import json
import os
import re
import socket
import struct
import subprocess
import sys
import threading
import time
import urllib.error
import urllib.request
from pathlib import Path

SWARM = "http://127.0.0.1:7801"
DATA = Path("/mnt/swarm_data")
SWARM_ROOT = Path(os.environ.get("SWARM_ROOT") or "/opt/swarmui")
DLBACKEND = DATA / "dlbackend"
COMFY_ROOT = DLBACKEND / "ComfyUI"
COMFY_VENV = COMFY_ROOT / "venv" / "bin" / "python"
SETTINGS = DATA / "Data" / "Settings.fds"
BACKENDS_FDS = DATA / "Data" / "Backends.fds"
# Comfy clone + torch can take 20–40+ min on a cold disk.
INSTALL_TIMEOUT = float(os.environ.get("GPU_RENT_COMFY_INSTALL_TIMEOUT") or "3600")


def fmt_bytes(n: float) -> str:
    n = float(n)
    for unit, div in (("GB", 1024**3), ("MB", 1024**2), ("KB", 1024), ("B", 1)):
        if n >= div or unit == "B":
            if unit == "B":
                return f"{int(n)}B"
            return f"{n / div:.1f}{unit}"
    return f"{n:.0f}B"


def progress_bar(done: int, total: int | None, *, width: int = 22) -> str:
    if not total or total <= 0:
        return ""
    pct = min(100.0, 100.0 * done / total)
    filled = min(width, max(0, int(width * done / total)))
    bar = "#" * filled + "-" * (width - filled)
    return f"[{bar}] {pct:5.1f}%  "


def dir_size_bytes(root: Path) -> int:
    if not root.is_dir():
        return 0
    try:
        out = subprocess.check_output(
            ["du", "-sb", str(root)],
            text=True,
            stderr=subprocess.DEVNULL,
            timeout=15,
        )
        return int(out.split()[0])
    except (subprocess.CalledProcessError, FileNotFoundError, OSError, ValueError, IndexError, subprocess.TimeoutExpired):
        pass
    total = 0
    try:
        for dirpath, _dirnames, filenames in os.walk(root):
            for name in filenames:
                try:
                    total += (Path(dirpath) / name).stat().st_size
                except OSError:
                    continue
    except OSError:
        return total
    return total


def detect_stage() -> str:
    """Best-effort stage while Linux comfy-install-linux.sh runs (no WS bytes)."""
    if not COMFY_ROOT.is_dir():
        return "ожидание clone ComfyUI"
    if not (COMFY_ROOT / ".git").is_dir() and not (COMFY_ROOT / "main.py").is_file():
        return "git clone ComfyUI…"
    if not COMFY_VENV.is_file():
        if (COMFY_ROOT / "venv").is_dir():
            return "создание venv…"
        return "ComfyUI clone ok → venv"
    # Prefer filesystem check — importing torch every second is too heavy.
    torch_dirs = list((COMFY_ROOT / "venv").glob("lib/python*/site-packages/torch"))
    if not torch_dirs:
        # pip may still be downloading into cache / tmp
        return "pip install torch (CUDA)…"
    req = COMFY_ROOT / "requirements.txt"
    if req.is_file():
        # Heuristic: if common deps missing, still on requirements.
        missing = False
        for pkg in ("aiohttp", "einops", "safetensors"):
            if not list((COMFY_ROOT / "venv").glob(f"lib/python*/site-packages/{pkg}*")):
                missing = True
                break
        if missing:
            return "pip install requirements…"
    return "Comfy deps почти готовы…"


def _journal_comfy_lines(since_sec: float = 90.0) -> list[str]:
    """Recent swarmui journal lines that look like comfy/pip progress."""
    try:
        out = subprocess.check_output(
            [
                "journalctl",
                "-u",
                "swarmui",
                "--no-pager",
                "-o",
                "cat",
                f"--since={int(max(1, since_sec))} seconds ago",
            ],
            text=True,
            stderr=subprocess.DEVNULL,
            timeout=8,
        )
    except (subprocess.CalledProcessError, FileNotFoundError, OSError, subprocess.TimeoutExpired):
        return []
    keys = (
        "comfy",
        "torch",
        "Downloading",
        "install",
        "Requirement",
        "Collecting",
        "Cloning",
        "Making venv",
        "Installation completed",
        "pip",
    )
    lines: list[str] = []
    for raw in out.splitlines():
        line = raw.strip()
        if not line or len(line) < 8:
            continue
        low = line.lower()
        if any(k.lower() in low for k in keys):
            # Strip noisy prefixes
            if "[ComfyUI Install" in line or "STDOUT" in line or "pip" in low or "torch" in low:
                lines.append(line[-180:])
    return lines[-6:]


def post(path: str, payload: dict, timeout: float = 15.0) -> dict:
    body = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        f"{SWARM}{path}",
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return json.loads(resp.read().decode("utf-8"))


def wait_http(deadline: float) -> None:
    last = ""
    while time.time() < deadline:
        try:
            req = urllib.request.Request(f"{SWARM}/", method="GET")
            with urllib.request.urlopen(req, timeout=5) as resp:
                if getattr(resp, "status", 200) == 200:
                    return
        except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError, OSError) as exc:
            last = str(exc)[:160]
        time.sleep(2)
    raise SystemExit(f"SwarmUI HTTP not up: {last}")


def get_session(deadline: float) -> str:
    last = ""
    while time.time() < deadline:
        try:
            data = post("/API/GetNewSession", {})
            sid = str(data.get("session_id") or "")
            if sid:
                return sid
            last = "no session_id"
        except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError, OSError, json.JSONDecodeError) as exc:
            last = str(exc)[:160]
        time.sleep(2)
    raise SystemExit(f"SwarmUI session unavailable: {last}")


def backend_status() -> str:
    try:
        sid = str(post("/API/GetNewSession", {}).get("session_id") or "")
        if not sid:
            return "unknown"
        data = post("/API/GetCurrentStatus", {"session_id": sid})
        be = data.get("backend_status") or {}
        return str(be.get("status") or "unknown").lower()
    except Exception as exc:
        return f"error:{exc}"


def backend_status_detail() -> tuple[str, str]:
    """Return (status, message) from GetCurrentStatus."""
    try:
        sid = str(post("/API/GetNewSession", {}).get("session_id") or "")
        if not sid:
            return "unknown", ""
        data = post("/API/GetCurrentStatus", {"session_id": sid})
        be = data.get("backend_status") or {}
        return (
            str(be.get("status") or "unknown").lower(),
            str(be.get("message") or "").strip(),
        )
    except Exception as exc:
        return f"error:{exc}", ""


def restart_backends(sid: str, *, which: str = "all") -> dict:
    """POST /API/RestartBackends — recovers ERRORED Comfy self-start."""
    return post(
        "/API/RestartBackends",
        {"session_id": sid, "backend": which},
        timeout=180.0,
    )


def recover_errored_backends(*, wait_sec: float = 180.0) -> str:
    """RestartBackends then poll until not errored (or timeout). Returns last status."""
    bstat, msg = backend_status_detail()
    print(
        f"backend errored — RestartBackends(all)"
        + (f" ({msg[:120]})" if msg else ""),
        flush=True,
    )
    sid = get_session(time.time() + 60)
    try:
        result = restart_backends(sid)
        print(f"RestartBackends: {result}", flush=True)
    except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError, OSError, json.JSONDecodeError) as exc:
        print(f"RestartBackends FAIL: {exc}", flush=True)
        return bstat

    deadline = time.time() + wait_sec
    last = "errored"
    while time.time() < deadline:
        time.sleep(8)
        last, msg2 = backend_status_detail()
        extra = f" — {msg2[:100]}" if msg2 else ""
        print(f"after restart: backend_status={last}{extra}", flush=True)
        if last != "errored" and not last.startswith("error:"):
            return last
    return last


def sanitize_backends_fds() -> bool:
    """Repair ExtraArgs: \\x --flag (FDS empty + appended flag). Return True if changed."""
    path = DATA / "Data" / "Backends.fds"
    if not path.is_file():
        return False
    try:
        text = path.read_text(encoding="utf-8")
    except OSError:
        return False
    new, n = re.subn(r"^(\s*ExtraArgs:\s*)\\x(\s+)", r"\1", text, flags=re.M)
    if not n:
        return False
    path.write_text(new, encoding="utf-8")
    print(f"sanitized Backends.fds ExtraArgs \\x corruption ({n} line(s))", flush=True)
    return True


def restart_swarmui_local() -> None:
    print("systemctl restart swarmui (после fix Backends.fds)", flush=True)
    try:
        subprocess.run(
            ["sudo", "-n", "systemctl", "restart", "swarmui"],
            check=False,
            timeout=120,
        )
    except (OSError, subprocess.TimeoutExpired) as exc:
        print(f"WARN restart swarmui: {exc}", flush=True)
    wait_http(time.time() + 180)


def list_backends(sid: str) -> dict:
    try:
        return post(
            "/API/ListBackends",
            {"session_id": sid, "nonreal": False, "full_data": False},
            timeout=30.0,
        )
    except Exception as exc:
        return {"error": str(exc)}


def add_comfy_selfstart(sid: str) -> dict:
    return post(
        "/API/AddNewBackend",
        {"session_id": sid, "type_id": "comfyui_selfstart"},
        timeout=60.0,
    )


def recover_empty_backends() -> str:
    """When API says empty but venv/IsInstalled exist — fix FDS and/or AddNewBackend."""
    changed = sanitize_backends_fds()
    if changed:
        restart_swarmui_local()
        bstat, _ = backend_status_detail()
        if bstat not in ("empty", "unknown") and not bstat.startswith("error:"):
            print(f"after sanitize: backend_status={bstat}", flush=True)
            return bstat

    sid = get_session(time.time() + 60)
    backends = list_backends(sid)
    if isinstance(backends, dict) and any(
        isinstance(v, dict) and v.get("type") for v in backends.values()
    ):
        bstat, _ = backend_status_detail()
        return bstat

    print("ListBackends empty — AddNewBackend comfyui_selfstart", flush=True)
    try:
        result = add_comfy_selfstart(sid)
        print(f"AddNewBackend: {result}", flush=True)
    except Exception as exc:
        print(f"AddNewBackend FAIL: {exc}", flush=True)
        return "empty"

    # Default StartScript may be dlbackend/comfy/...; our tree is dlbackend/ComfyUI.
    time.sleep(2)
    bstat, _ = backend_status_detail()
    return bstat


def run_diagnostics() -> None:
    """Best-effort: prefer uploaded swarm_diag.py, else journalctl snippet."""
    diag = Path("/tmp/gpu-rent-swarm_diag.py")
    if diag.is_file():
        try:
            subprocess.run(
                [sys.executable, str(diag)],
                check=False,
                timeout=90,
            )
            return
        except (OSError, subprocess.TimeoutExpired) as exc:
            print(f"diag script fail: {exc}", flush=True)
    print("=== fallback journalctl -u swarmui -n 60 ===", flush=True)
    try:
        out = subprocess.check_output(
            ["journalctl", "-u", "swarmui", "-n", "60", "--no-pager"],
            text=True,
            stderr=subprocess.STDOUT,
            timeout=30,
            errors="replace",
        )
        print(out, flush=True)
    except (OSError, subprocess.CalledProcessError, subprocess.TimeoutExpired) as exc:
        print(f"journalctl fail: {exc}", flush=True)


def settings_is_installed() -> bool | None:
    if not SETTINGS.is_file():
        return None
    try:
        text = SETTINGS.read_text(encoding="utf-8", errors="replace")
    except OSError:
        return None
    for line in text.splitlines():
        if "IsInstalled" in line:
            low = line.lower()
            if "true" in low:
                return True
            if "false" in low:
                return False
    return None


def comfy_venv_ok() -> bool:
    return COMFY_VENV.is_file() and os.access(COMFY_VENV, os.X_OK)


def _recv_exact(sock: socket.socket, n: int) -> bytes:
    buf = b""
    while len(buf) < n:
        chunk = sock.recv(n - len(buf))
        if not chunk:
            raise ConnectionError("websocket closed")
        buf += chunk
    return buf


def _ws_handshake(host: str, port: int, path: str) -> socket.socket:
    key = base64.b64encode(os.urandom(16)).decode("ascii")
    req = (
        f"GET {path} HTTP/1.1\r\n"
        f"Host: {host}:{port}\r\n"
        "Upgrade: websocket\r\n"
        "Connection: Upgrade\r\n"
        f"Sec-WebSocket-Key: {key}\r\n"
        "Sec-WebSocket-Version: 13\r\n"
        "\r\n"
    ).encode("ascii")
    sock = socket.create_connection((host, port), timeout=30)
    sock.sendall(req)
    # Read headers until blank line
    data = b""
    while b"\r\n\r\n" not in data:
        chunk = sock.recv(4096)
        if not chunk:
            sock.close()
            raise ConnectionError("no WS handshake response")
        data += chunk
        if len(data) > 65536:
            sock.close()
            raise ConnectionError("WS handshake too large")
    head = data.split(b"\r\n\r\n", 1)[0].decode("latin-1", errors="replace")
    if "101" not in head.split("\r\n", 1)[0]:
        sock.close()
        raise ConnectionError(f"WS handshake failed: {head.splitlines()[0]}")
    expect = base64.b64encode(
        hashlib.sha1(
            (key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode("ascii")
        ).digest()
    ).decode("ascii")
    if expect not in head:
        # Some servers omit echoing exact key in proxies; 101 is enough.
        pass
    sock.settimeout(120.0)
    return sock


def _ws_send_text(sock: socket.socket, text: str) -> None:
    payload = text.encode("utf-8")
    mask = os.urandom(4)
    header = bytearray([0x81])  # FIN + text
    n = len(payload)
    if n < 126:
        header.append(0x80 | n)
    elif n < 65536:
        header.append(0x80 | 126)
        header.extend(struct.pack("!H", n))
    else:
        header.append(0x80 | 127)
        header.extend(struct.pack("!Q", n))
    header.extend(mask)
    masked = bytes(b ^ mask[i % 4] for i, b in enumerate(payload))
    sock.sendall(header + masked)


def _ws_recv_text(sock: socket.socket) -> str | None:
    """Return next text frame payload, or None on close."""
    while True:
        hdr = _recv_exact(sock, 2)
        opcode = hdr[0] & 0x0F
        masked = bool(hdr[1] & 0x80)
        length = hdr[1] & 0x7F
        if length == 126:
            length = struct.unpack("!H", _recv_exact(sock, 2))[0]
        elif length == 127:
            length = struct.unpack("!Q", _recv_exact(sock, 8))[0]
        mask_key = _recv_exact(sock, 4) if masked else b""
        raw = _recv_exact(sock, length) if length else b""
        if masked:
            raw = bytes(b ^ mask_key[i % 4] for i, b in enumerate(raw))
        if opcode == 0x8:  # close
            return None
        if opcode == 0x9:  # ping → pong
            # build pong
            frame = bytearray([0x8A, 0x80 | len(raw)])
            m = os.urandom(4)
            frame.extend(m)
            frame.extend(bytes(b ^ m[i % 4] for i, b in enumerate(raw)))
            sock.sendall(frame)
            continue
        if opcode in (0x1, 0x0):  # text / continuation
            return raw.decode("utf-8", errors="replace")
        # ignore binary / other


def run_install(sid: str) -> None:
    print("SwarmUI InstallConfirmWS: backend=comfyui models=none …", flush=True)
    sock = _ws_handshake("127.0.0.1", 7801, "/API/InstallConfirmWS")
    payload = {
        "session_id": sid,
        "theme": "modern_dark",
        "installed_for": "just_self",
        "backend": "comfyui",
        "models": "none",
        "install_amd": False,
        "language": "en",
        "make_shortcut": False,
    }
    _ws_send_text(sock, json.dumps(payload))
    deadline = time.time() + INSTALL_TIMEOUT
    last_info = ""
    last_journal: set[str] = set()
    stop = threading.Event()
    t0 = time.monotonic()
    size0 = dir_size_bytes(DLBACKEND)
    last_size = size0
    last_size_t = t0
    # WS byte progress (Windows 7z / model downloads); Linux script usually silent.
    ws_prog: dict = {"done": 0, "total": 0, "steps": 0, "total_steps": 0, "t": 0.0}

    def emit_progress_line(*, final: bool = False) -> None:
        nonlocal last_size, last_size_t
        now = time.monotonic()
        elapsed = max(now - t0, 0.001)
        size = dir_size_bytes(DLBACKEND)
        dt = max(now - last_size_t, 0.001)
        speed = (size - last_size) / dt if size >= last_size else 0.0
        last_size, last_size_t = size, now
        stage = detect_stage()
        grown = max(0, size - size0)
        parts = [
            f"[installer] {stage}",
            f"dlbackend={fmt_bytes(size)}",
            f"+{fmt_bytes(grown)}",
            f"{fmt_bytes(speed)}/s",
            f"{int(elapsed)}s",
        ]
        done = int(ws_prog.get("done") or 0)
        total = int(ws_prog.get("total") or 0)
        if done > 0 or total > 0:
            bar = progress_bar(done, total if total > 0 else None)
            steps = ws_prog.get("steps")
            total_steps = ws_prog.get("total_steps")
            parts.insert(
                1,
                f"step {steps}/{total_steps} {bar}{fmt_bytes(done)}"
                + (f"/{fmt_bytes(total)}" if total else ""),
            )
        line = " · ".join(parts)
        if final:
            print(line, flush=True)
        else:
            print(line, end="\r", flush=True)

    def monitor() -> None:
        while not stop.wait(1.0):
            emit_progress_line()
            # Surface interesting journal lines as real log lines (newline).
            for jline in _journal_comfy_lines(120.0):
                if jline in last_journal:
                    continue
                last_journal.add(jline)
                # End in-place progress then print journal snippet.
                print(flush=True)
                print(f"  … {jline}", flush=True)

    mon = threading.Thread(target=monitor, name="comfy-install-progress", daemon=True)
    mon.start()
    try:
        while time.time() < deadline:
            try:
                sock.settimeout(max(5.0, min(30.0, deadline - time.time())))
                msg = _ws_recv_text(sock)
            except (TimeoutError, socket.timeout):
                continue
            if msg is None:
                raise SystemExit("InstallConfirmWS closed without success")
            try:
                data = json.loads(msg)
            except json.JSONDecodeError:
                print(f"ws raw: {msg[:200]}", flush=True)
                continue
            if data.get("error"):
                raise SystemExit(f"InstallConfirmWS error: {data['error']}")
            if data.get("info"):
                info = str(data["info"])
                if info != last_info:
                    print(flush=True)  # break progress \r line
                    print(f"[installer] {info}", flush=True)
                    last_info = info
            if "progress" in data:
                try:
                    ws_prog["done"] = int(data.get("progress") or 0)
                    ws_prog["total"] = int(data.get("total") or 0)
                    ws_prog["steps"] = int(data.get("steps") or 0)
                    ws_prog["total_steps"] = int(data.get("total_steps") or 0)
                    ws_prog["t"] = time.monotonic()
                except (TypeError, ValueError):
                    pass
                # Byte tick from Swarm (rare on Linux) — refresh bar now.
                if ws_prog["done"] or ws_prog["total"]:
                    emit_progress_line()
            if data.get("success"):
                stop.set()
                print(flush=True)
                emit_progress_line(final=True)
                print("InstallConfirmWS: success", flush=True)
                return
        raise SystemExit(f"InstallConfirmWS timeout after {int(INSTALL_TIMEOUT)}s")
    finally:
        stop.set()
        try:
            sock.close()
        except OSError:
            pass
        mon.join(timeout=2.0)


def main() -> int:
    # WorkingDirectory for comfy-install-linux.sh is Swarm root when launched by
    # Installation.cs; InstallConfirmWS handles that for us.
    wait_http(time.time() + 180)
    bstat, bmsg = backend_status_detail()
    print(
        f"backend_status={bstat} venv={'yes' if comfy_venv_ok() else 'no'} "
        f"IsInstalled={settings_is_installed()}"
        + (f" msg={bmsg[:160]}" if bmsg else ""),
        flush=True,
    )

    # Errored backends are registered but dead — restart, don't skip as "present".
    if bstat == "errored":
        bstat = recover_errored_backends(wait_sec=180.0)
        if bstat == "errored" or bstat.startswith("error:"):
            print(
                "FAIL: backend still errored after RestartBackends — collecting diag",
                file=sys.stderr,
            )
            run_diagnostics()
            return 1
        if bstat in ("running", "idle", "loading", "some_loading") and comfy_venv_ok():
            print(f"recovered to {bstat} — skip InstallConfirmWS")
            return 0

    # Backends already registered (healthy / loading / suspended) + venv → skip.
    if bstat == "idle":
        print("backends present (idle/suspended) + skip install")
        return 0
    if bstat not in ("empty", "unknown", "errored") and not bstat.startswith("error:"):
        if comfy_venv_ok():
            print(f"backends present ({bstat}) + venv — skip install")
            return 0

    installed = settings_is_installed()
    if installed is True and not comfy_venv_ok() and bstat == "empty":
        print(
            "WARN: Settings IsInstalled=true but backends empty and no Comfy venv. "
            "Open SwarmUI → Server → Backends and add ComfyUI Self-Starting, "
            "or delete Data/Settings.fds IsInstalled and re-run up."
        )
        run_diagnostics()
        return 1

    if installed is True and comfy_venv_ok() and bstat not in ("empty", "errored"):
        print("IsInstalled + venv — skip InstallConfirmWS (ждём ready отдельно)")
        return 0

    if comfy_venv_ok() and bstat == "empty":
        print(
            "backend empty при venv/IsInstalled — recover (sanitize FDS / AddNewBackend)…",
            flush=True,
        )
        bstat = recover_empty_backends()
        if bstat not in ("empty", "unknown") and not bstat.startswith("error:"):
            print(f"recovered empty → {bstat}")
            return 0
        if installed is True:
            print(
                "WARN: всё ещё empty после recover — пробую InstallConfirmWS",
                flush=True,
            )
        else:
            print("venv есть, IsInstalled=false — запускаю InstallConfirmWS")

    marker = DATA / ".gpu-rent-comfy-installing"
    try:
        marker.write_text(f"{int(time.time())}\n", encoding="utf-8")
    except OSError:
        pass
    try:
        sid = get_session(time.time() + 120)
        run_install(sid)
    finally:
        try:
            marker.unlink(missing_ok=True)
        except OSError:
            pass

    # Confirm outcome
    time.sleep(3)
    bstat2 = backend_status()
    print(
        f"after install: backend_status={bstat2} "
        f"venv={'yes' if comfy_venv_ok() else 'no'}"
    )
    if bstat2 == "errored":
        bstat2 = recover_errored_backends(wait_sec=120.0)
        if bstat2 == "errored":
            print("FAIL: install done but backend errored", file=sys.stderr)
            run_diagnostics()
            return 1
    if not comfy_venv_ok() and bstat2 == "empty":
        print("FAIL: install finished but still empty / no venv", file=sys.stderr)
        run_diagnostics()
        return 1
    return 0


if __name__ == "__main__":
    # Ensure Installation.cs cwd paths resolve if anything shells out relative —
    # InstallConfirmWS itself cds via WorkingDirectory of the service.
    try:
        os.chdir(str(SWARM_ROOT))
    except OSError:
        pass
    raise SystemExit(main())