#!/usr/bin/env bash
# download_missing_models.sh
# Interactively pick workflows from ./workflows, then:
#   1. Find every custom-node pack the workflows depend on (via each node's
#      cnr_id/aux_id metadata, with Comfy Registry + ComfyUI-Manager DB lookups
#      as fallbacks), git-clone them into ComfyUI/custom_nodes, and install
#      their Python requirements — so they're live on the next ComfyUI restart.
#   2. Find every model they reference, locate those models in the HF repo
#      aleph65/ComfyUI (which mirrors ComfyUI/models/), and download the
#      missing ones into place, in parallel, with progress.
#
# Usage:
#   ./download_missing_models.sh                 # interactive workflow picker
#   ./download_missing_models.sh --all           # select every workflow
#   ./download_missing_models.sh qwen-edit.json  # select specific workflow(s)
#   ./download_missing_models.sh --dry-run ...   # show the plan, install/download nothing
#   ./download_missing_models.sh --skip-nodes .. # models only, skip custom-node installs
#   ./download_missing_models.sh --yes ...       # skip the "Proceed?" confirmation
#   ./download_missing_models.sh --no-refresh .. # don't re-sync ./workflows from the HF repo first
#
# Auth: uses the HUGGING_FACE_ACCESS_TOKEN environment variable if set
# (falls back to HF_TOKEN), otherwise prompts for a token.

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
export HF_HUB_ENABLE_HF_TRANSFER=1
# never let git sit on an interactive credential/hostkey prompt
export GIT_TERMINAL_PROMPT=0

# 1) Hugging Face token: HUGGING_FACE_ACCESS_TOKEN, then HF_TOKEN, then prompt.
if [[ -n "${HUGGING_FACE_ACCESS_TOKEN:-}" ]]; then
    HF_TOKEN="$HUGGING_FACE_ACCESS_TOKEN"
    echo "Using HF token from environment (HUGGING_FACE_ACCESS_TOKEN)."
elif [[ -n "${HF_TOKEN:-}" ]]; then
    echo "Using HF token from environment (HF_TOKEN)."
else
    read -r -s -p "Enter your Hugging Face token (input hidden): " HF_TOKEN
    echo
    [[ -n "$HF_TOKEN" ]] || { echo "No token provided — aborting."; exit 1; }
fi
export HF_TOKEN

# 2) Ensure deps (huggingface_hub >= 1.x ships fast Xet downloads built in and no
# longer provides the hf-transfer extra, so hf_transfer is best-effort only)
if ! python3 -c "import huggingface_hub" >/dev/null 2>&1; then
    echo "Installing huggingface_hub[hf_transfer]..."
    pip install -q -U "huggingface_hub[hf_transfer]" || pip install -q -U huggingface_hub
fi
python3 -c "import hf_transfer" >/dev/null 2>&1 || pip install -q hf_transfer >/dev/null 2>&1 || true

# 3) Run the downloader.
# Load the Python code into a variable (NOT via stdin, which must stay attached
# to the terminal for the interactive prompts).
PYCODE=$(cat <<'PYEOF'
import json
import os
import re
import socket
import subprocess
import sys
import glob
import threading
import time
import urllib.request

# belt-and-braces: no bare socket (HF API pagination etc.) may stall forever
socket.setdefaulttimeout(60)

WORKSPACE = os.environ.get("COMFY_WORKSPACE", "/workspace")
WORKFLOWS_DIR = os.path.join(WORKSPACE, "workflows")
COMFY_DIR = os.path.join(WORKSPACE, "ComfyUI")
CUSTOM_NODES_DIR = os.path.join(COMFY_DIR, "custom_nodes")
REPO_ID = "aleph65/ComfyUI"
MODEL_EXTS = (".safetensors", ".ckpt", ".pt", ".pth", ".bin", ".gguf", ".sft", ".onnx")
SKIP_NODE_TYPES = {"MarkdownNote", "Note", "PrimitiveString", "String"}
# node types that never imply a custom-node pack
DEP_SKIP_TYPES = {"MarkdownNote", "Note", "Reroute", "PrimitiveNode"}
REGISTRY_URL = "https://api.comfy.org/nodes/{}"
NODE_MAP_URL = "https://raw.githubusercontent.com/ltdrdata/ComfyUI-Manager/main/extension-node-map.json"

C_RESET, C_GREEN, C_YELLOW, C_RED, C_CYAN, C_BOLD = "\033[0m", "\033[32m", "\033[33m", "\033[31m", "\033[36m", "\033[1m"

def human(n):
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if n < 1024 or unit == "TB":
            return f"{n:.1f} {unit}" if unit != "B" else f"{int(n)} B"
        n /= 1024

# ---------------- workflow selection ----------------

def pick_workflows(files, argv):
    names = [os.path.basename(f) for f in files]
    args = [a for a in argv if not a.startswith("-")]
    if "--all" in argv:
        return files
    if args:
        chosen = []
        for a in args:
            matches = [f for f in files if os.path.basename(f) == a or os.path.basename(f) == a + ".json"]
            if not matches:
                sys.exit(f"No workflow named '{a}' in {WORKFLOWS_DIR}")
            chosen += matches
        return chosen
    print(f"\n{C_BOLD}Available workflows:{C_RESET}")
    for i, name in enumerate(names, 1):
        print(f"  {i}) {name}")
    print()
    while True:
        try:
            raw = input("Select workflows (e.g. '1 3', '1,3', or 'all'; q to quit): ").strip().lower()
        except EOFError:
            sys.exit("\nNo input available. Run interactively, or pass workflow names or --all.")
        if raw in ("q", "quit", "exit"):
            sys.exit("Aborted.")
        if raw in ("all", "a", "*"):
            return files
        picks = raw.replace(",", " ").split()
        if picks and all(p.isdigit() and 1 <= int(p) <= len(names) for p in picks):
            return [files[int(p) - 1] for p in dict.fromkeys(picks)]
        print(f"  Invalid selection — enter numbers 1-{len(names)}, or 'all'.")

# ---------------- model extraction ----------------

def looks_like_model(v):
    if not isinstance(v, str):
        return False
    v = v.strip()
    if not v.lower().endswith(MODEL_EXTS):
        return False
    if "\n" in v or "http://" in v or "https://" in v or len(v) > 300:
        return False
    return True

def iter_graph_nodes(data):
    """Yield every node dict in a workflow, including subgraph instances/definitions."""
    def walk(nodes):
        for n in nodes:
            if not isinstance(n, dict):
                continue
            yield n
            sub = n.get("subgraph")
            if isinstance(sub, dict) and isinstance(sub.get("nodes"), list):
                yield from walk(sub["nodes"])
    if isinstance(data, dict) and isinstance(data.get("nodes"), list):
        yield from walk(data["nodes"])
        for sg in (data.get("definitions") or {}).get("subgraphs", []) or []:
            if isinstance(sg, dict) and isinstance(sg.get("nodes"), list):
                yield from walk(sg["nodes"])
    elif isinstance(data, dict):
        # API format: {"1": {"class_type": ..., "inputs": {...}}, ...}
        for n in data.values():
            if isinstance(n, dict) and "class_type" in n:
                yield n

def extract_models(data):
    """Return set of model refs (may include subdirs like 'flux2/klein/x.safetensors')."""
    refs = set()

    def scan_values(values):
        stack = [values]
        while stack:
            v = stack.pop()
            if isinstance(v, str) and looks_like_model(v):
                refs.add(v.strip().replace("\\", "/"))
            elif isinstance(v, list):
                stack.extend(v)
            elif isinstance(v, dict):
                stack.extend(v.values())

    for n in iter_graph_nodes(data):
        if n.get("type") in SKIP_NODE_TYPES or n.get("class_type") in SKIP_NODE_TYPES:
            continue
        if "widgets_values" in n:
            scan_values(n["widgets_values"])
        if isinstance(n.get("inputs"), dict):
            scan_values(list(n["inputs"].values()))
    return refs

# ---------------- custom-node dependency extraction ----------------

def norm_name(s):
    return re.sub(r"[-_.\s]", "", (s or "").lower())

def fetch_json(url, timeout=30, deadline=None):
    """GET a JSON URL. `timeout` bounds each socket op; `deadline` bounds the
    whole download (a slow-dripping CDN can otherwise stretch a 2 MB fetch
    to many silent minutes without ever tripping the per-read timeout)."""
    req = urllib.request.Request(url, headers={"User-Agent": "comfy-dep-installer"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        if deadline is None:
            return json.load(r)
        start, chunks = time.time(), []
        while True:
            if time.time() - start > deadline:
                raise TimeoutError(f"download exceeded {deadline}s")
            chunk = r.read(65536)
            if not chunk:
                break
            chunks.append(chunk)
        return json.loads(b"".join(chunks))

_node_map_cache = None
def node_map():
    """ComfyUI-Manager's extension-node-map: repo URL -> list of node types.
    Cached on disk for a day — the fetch can crawl on pod networks."""
    global _node_map_cache
    if _node_map_cache is None:
        cache_file = os.path.join(WORKSPACE, ".cache", "extension-node-map.json")
        try:
            if os.path.exists(cache_file) and time.time() - os.path.getmtime(cache_file) < 86400:
                with open(cache_file) as fh:
                    raw = json.load(fh)
            else:
                print(f"  {C_CYAN}fetching ComfyUI-Manager node map (~2 MB — can take a minute on slow networks) ...{C_RESET}")
                raw = fetch_json(NODE_MAP_URL, deadline=180)
                os.makedirs(os.path.dirname(cache_file), exist_ok=True)
                with open(cache_file, "w") as fh:
                    json.dump(raw, fh)
            _node_map_cache = {url: entry[0] for url, entry in raw.items()
                               if isinstance(entry, list) and entry and isinstance(entry[0], list)}
        except Exception as e:
            print(f"  {C_YELLOW}⚠ could not fetch ComfyUI-Manager node map ({e}){C_RESET}")
            _node_map_cache = {}
    return _node_map_cache

def extract_node_deps(data):
    """Return ({pack_key: dep}, {propertyless node types}).
    dep = {'cnr_id', 'aux_id', 'ver', 'types': set}"""
    deps, orphan_types = {}, set()
    uuid_re = re.compile(r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$", re.I)
    for n in iter_graph_nodes(data):
        ntype = n.get("type") or n.get("class_type") or ""
        if ntype in DEP_SKIP_TYPES or uuid_re.match(ntype):
            continue
        props = n.get("properties") or {}
        cnr_id, aux_id, ver = props.get("cnr_id"), props.get("aux_id"), props.get("ver")
        if cnr_id == "comfy-core":
            continue
        if cnr_id or aux_id:
            key = norm_name(cnr_id or aux_id.split("/")[-1])
            dep = deps.setdefault(key, {"cnr_id": cnr_id, "aux_id": aux_id, "ver": ver, "types": set()})
            dep["types"].add(ntype)
            # a commit-hash ver is more specific than semver; keep the hash if seen
            if ver and re.fullmatch(r"[0-9a-f]{40}", str(dep.get("ver") or ""), re.I) is None:
                dep["ver"] = ver
        elif ntype:
            orphan_types.add(ntype)
    return deps, orphan_types

def grep_defines(ntype, paths):
    """True if any python file under paths mentions the node type as a whole word."""
    targets = [p for p in paths if os.path.exists(p)]
    if not targets:
        return False
    r = subprocess.run(["grep", "-rlwF", "--include=*.py", "-e", ntype] + targets,
                       stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    return r.returncode == 0

def is_core_node(ntype):
    return grep_defines(ntype, [
        os.path.join(COMFY_DIR, "nodes.py"),
        os.path.join(COMFY_DIR, "comfy_extras"),
        os.path.join(COMFY_DIR, "comfy_api_nodes"),
    ])

def installed_packs():
    """norm_name(dir) -> dir for every pack already in custom_nodes."""
    out = {}
    if os.path.isdir(CUSTOM_NODES_DIR):
        for d in os.listdir(CUSTOM_NODES_DIR):
            full = os.path.join(CUSTOM_NODES_DIR, d)
            if os.path.isdir(full) and not d.startswith((".", "__")):
                out[norm_name(d)] = d
    return out

def provided_by_installed(ntype):
    """True if an installed pack plausibly defines this node type. Display
    names like 'Power Lora Loader (rgthree)' are often built dynamically, so
    the exact string never appears in .py source — also try the base name and
    match a parenthetical suffix against installed pack dir names."""
    if grep_defines(ntype, [CUSTOM_NODES_DIR]):
        return True
    m = re.match(r"^(.+?)\s*\(([^()]+)\)$", ntype)
    if m:
        base, suffix = m.group(1), norm_name(m.group(2))
        if suffix and any(suffix in key for key in installed_packs()):
            return True
        if grep_defines(base, [CUSTOM_NODES_DIR]):
            return True
    return False

def resolve_repo_url(dep):
    """Find the git URL for a dep. Registry first, then aux_id, then Manager DB."""
    if dep.get("cnr_id"):
        try:
            info = fetch_json(REGISTRY_URL.format(dep["cnr_id"]))
            if info.get("repository"):
                return info["repository"], "comfy registry"
        except Exception:
            pass
    if dep.get("aux_id") and "/" in dep["aux_id"]:
        return f"https://github.com/{dep['aux_id']}", "aux_id"
    # fall back to ComfyUI-Manager's node map: which repos provide these node types?
    want_key = norm_name(dep.get("cnr_id") or "")
    candidates = [url for url, types in node_map().items()
                  if dep["types"] & set(types)]
    if candidates:
        # prefer a repo whose name matches the cnr_id
        exact = [u for u in candidates if norm_name(u.rstrip("/").split("/")[-1]) == want_key]
        return (exact or candidates)[0], "ComfyUI-Manager node map"
    return None, None

GIT_ENV = {**os.environ, "GIT_TERMINAL_PROMPT": "0"}

def install_pack(name, url, ver):
    dest = os.path.join(CUSTOM_NODES_DIR, url.rstrip("/").split("/")[-1].removesuffix(".git"))
    print(f"\n{C_CYAN}Installing {name} from {url} ...{C_RESET}")
    if not os.path.isdir(dest):
        subprocess.run(["git", "clone", "--recursive", "--progress", url, dest],
                       check=True, env=GIT_ENV, timeout=900, stdin=subprocess.DEVNULL)
    if ver and re.fullmatch(r"[0-9a-f]{40}", str(ver), re.I):
        r = subprocess.run(["git", "-C", dest, "checkout", ver], env=GIT_ENV, timeout=120,
                           stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        print(f"  pinned to commit {ver[:12]}" if r.returncode == 0
              else f"  {C_YELLOW}⚠ could not pin to {ver[:12]}, staying on latest{C_RESET}")
    req = os.path.join(dest, "requirements.txt")
    if os.path.exists(req):
        print("  installing python requirements (can take a few minutes) ...")
        r = subprocess.run([sys.executable, "-m", "pip", "install", "-r", req,
                            "--progress-bar", "off"],
                           timeout=1800, stdin=subprocess.DEVNULL,
                           stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
        if r.returncode != 0:
            print(r.stdout)
            print(f"  {C_YELLOW}⚠ pip install had errors — see output above{C_RESET}")
        else:
            installed = [l for l in r.stdout.splitlines() if l.startswith("Successfully installed")]
            print(f"  {installed[0] if installed else 'requirements already satisfied'}")
    inst = os.path.join(dest, "install.py")
    if os.path.exists(inst):
        print("  running install.py ...")
        r = subprocess.run([sys.executable, inst], cwd=dest, timeout=1800,
                           stdin=subprocess.DEVNULL)
        if r.returncode != 0:
            print(f"  {C_YELLOW}⚠ install.py exited with {r.returncode}{C_RESET}")
    return dest

def plan_node_deps(chosen_data, dry_run):
    """Collect deps across workflows, resolve, and return list of (name, url, ver) to install."""
    all_deps, all_orphans = {}, set()
    for wf_name, data in chosen_data:
        deps, orphans = extract_node_deps(data)
        for k, d in deps.items():
            tgt = all_deps.setdefault(k, d)
            if tgt is not d:
                tgt["types"] |= d["types"]
                # prefer a commit-hash pin over a semver, else keep the first seen
                if d["ver"] and not re.fullmatch(r"[0-9a-f]{40}", str(tgt["ver"] or ""), re.I):
                    tgt["ver"] = d["ver"]
        all_orphans |= orphans

    installed = installed_packs()

    # nodes with no pack metadata: fine if core or provided by an installed pack,
    # otherwise try to identify the pack via the Manager node map
    unresolved_types = []
    for t in sorted(all_orphans):
        if any(t in d["types"] for d in all_deps.values()):
            continue
        if is_core_node(t) or provided_by_installed(t):
            continue
        hits = [url for url, types in node_map().items() if t in types]
        if hits:
            key = norm_name(hits[0].rstrip("/").split("/")[-1])
            dep = all_deps.setdefault(key, {"cnr_id": None, "aux_id": None, "ver": None, "types": set()})
            dep["types"].add(t)
            dep["_url"] = hits[0]
        else:
            unresolved_types.append(t)

    to_install, already, failed = [], [], []
    for key, dep in sorted(all_deps.items()):
        name = dep.get("cnr_id") or dep.get("aux_id") or key
        if key in installed:
            already.append((name, installed[key]))
            continue
        url, source = (dep.get("_url"), "ComfyUI-Manager node map") if dep.get("_url") else resolve_repo_url(dep)
        if url:
            repo_dir = norm_name(url.rstrip("/").split("/")[-1].removesuffix(".git"))
            if repo_dir in installed:
                already.append((name, installed[repo_dir]))
            else:
                to_install.append((name, url, dep.get("ver"), source))
        else:
            failed.append((name, dep["types"]))

    print(f"\n{C_BOLD}Custom-node dependencies:{C_RESET}")
    if not (to_install or already or failed):
        print(f"  {C_GREEN}✔ none needed — selected workflows only use core nodes.{C_RESET}")
    for name, d in already:
        print(f"  {C_GREEN}✔ installed{C_RESET}  {name}  (custom_nodes/{d})")
    for name, url, ver, source in to_install:
        pin = f", pin {str(ver)[:12]}" if ver and re.fullmatch(r"[0-9a-f]{40}", str(ver), re.I) else ""
        print(f"  {C_YELLOW}↓ install{C_RESET}    {name}  ({url}{pin}, via {source})")
    for name, types in failed:
        print(f"  {C_RED}✘ unknown{C_RESET}    {name}  — could not find a repo (nodes: {', '.join(sorted(types))})")
    for t in unresolved_types:
        print(f"  {C_RED}✘ unknown{C_RESET}    node type '{t}' — not core, not installed, not in the Manager DB")
    return to_install

# ---------------- main ----------------

def cleanup_debris():
    """Remove partial-download fragments and stale locks left by interrupted runs."""
    cache = os.path.join(COMFY_DIR, ".cache", "huggingface", "download")
    freed, count = 0, 0
    for root, _, names in os.walk(cache):
        for name in names:
            if name.endswith((".incomplete", ".lock")):
                p = os.path.join(root, name)
                try:
                    freed += os.path.getsize(p)
                    os.remove(p)
                    count += 1
                except OSError:
                    pass
    if count:
        print(f"{C_YELLOW}Cleaned up {count} leftover partial-download file(s) ({human(freed)} freed).{C_RESET}")

def download_heartbeat(to_download):
    """Print bytes-on-disk every 20s so a big download never looks hung."""
    total = sum(s for _, _, s, _ in to_download)
    stop = threading.Event()
    cache = os.path.join(COMFY_DIR, ".cache", "huggingface", "download")
    def run():
        while not stop.wait(20):
            done = 0
            for _, _, size, dest in to_download:
                if os.path.exists(dest):
                    done += min(os.path.getsize(dest), size)
            for root, _, names in os.walk(cache):
                for n in names:
                    if n.endswith(".incomplete"):
                        try:
                            done += os.path.getsize(os.path.join(root, n))
                        except OSError:
                            pass
            print(f"  ... still downloading: ~{human(min(done, total))} of {human(total)} on disk", flush=True)
    t = threading.Thread(target=run, daemon=True)
    t.start()
    return stop

def main():
    argv = sys.argv[1:]
    dry_run = "--dry-run" in argv
    skip_nodes = "--skip-nodes" in argv
    assume_yes = "--yes" in argv or "-y" in argv

    cleanup_debris()

    if not os.path.isdir(WORKFLOWS_DIR):
        print(f"{C_YELLOW}{WORKFLOWS_DIR} not found — downloading workflows from hf.co/{REPO_ID} ...{C_RESET}")
        from huggingface_hub import snapshot_download
        snapshot_download(
            repo_id=REPO_ID,
            allow_patterns=["workflows/*.json"],
            local_dir=WORKSPACE,
            token=os.environ.get("HF_TOKEN"),
        )
    elif "--no-refresh" not in argv:
        # The workflows baked into the pod image can lag the HF repo, and a
        # stale copy makes the model plan silently miss anything added to a
        # workflow since the image was built. Sync workflows/*.json from the
        # repo before planning; the repo wins over local edits, so pass
        # --no-refresh to plan against the local copies instead.
        print(f"{C_CYAN}Syncing workflows from hf.co/{REPO_ID} (--no-refresh to skip) ...{C_RESET}")
        try:
            from huggingface_hub import snapshot_download
            snapshot_download(
                repo_id=REPO_ID,
                allow_patterns=["workflows/*.json"],
                local_dir=WORKSPACE,
                token=os.environ.get("HF_TOKEN"),
            )
        except Exception as e:
            print(f"{C_YELLOW}⚠ workflow sync failed ({e}) — planning against the local copies.{C_RESET}")

    files = sorted(glob.glob(os.path.join(WORKFLOWS_DIR, "*.json")))
    if not files:
        sys.exit(f"No workflow JSONs found in {WORKFLOWS_DIR}")

    chosen = pick_workflows(files, argv)
    chosen_data = []
    print(f"\n{C_BOLD}Selected workflows:{C_RESET}")
    wanted = {}  # ref -> [workflow names]
    for f in chosen:
        with open(f) as fh:
            data = json.load(fh)
        chosen_data.append((os.path.basename(f), data))
        refs = extract_models(data)
        print(f"  • {os.path.basename(f)}  ({len(refs)} model refs)")
        for r in refs:
            wanted.setdefault(r, []).append(os.path.basename(f))

    # -------- custom-node packs --------
    node_installs = [] if skip_nodes else plan_node_deps(chosen_data, dry_run)

    # -------- models --------
    to_download, have, not_found = [], [], []
    if wanted:
        print(f"\n{C_CYAN}Listing files in hf.co/{REPO_ID} ...{C_RESET}")
        from huggingface_hub import HfApi
        api = HfApi(token=os.environ.get("HF_TOKEN"))
        repo_files = {}  # repo path -> size
        for entry in api.list_repo_tree(REPO_ID, recursive=True):
            if hasattr(entry, "size") and entry.path.startswith("models/"):
                repo_files[entry.path] = entry.size or 0

        # index by basename and by suffix path
        by_basename = {}
        for p in repo_files:
            by_basename.setdefault(os.path.basename(p), []).append(p)

        for ref, wfs in sorted(wanted.items()):
            base = os.path.basename(ref)
            candidates = by_basename.get(base, [])
            # prefer a repo path that ends with the workflow's relative path (handles lora subdirs)
            match = next((p for p in candidates if p.endswith("/" + ref) or p == "models/" + ref), None)
            if match is None and candidates:
                match = candidates[0]
            if match is None:
                not_found.append((ref, wfs))
                continue
            dest = os.path.join(COMFY_DIR, match)
            size = repo_files[match]
            if os.path.exists(dest) and (size == 0 or os.path.getsize(dest) == size):
                have.append((ref, match, size))
            else:
                to_download.append((ref, match, size, dest))

        print(f"\n{C_BOLD}Model plan:{C_RESET}")
        for ref, match, size in have:
            print(f"  {C_GREEN}✔ have{C_RESET}     {match}  ({human(size)})")
        for ref, match, size, dest in to_download:
            print(f"  {C_YELLOW}↓ fetch{C_RESET}    {match}  ({human(size)})")
        for ref, wfs in not_found:
            print(f"  {C_RED}✘ missing{C_RESET}  {ref}  — not in hf.co/{REPO_ID} (used by {', '.join(wfs)})")

        total = sum(s for _, _, s, _ in to_download)
        print(f"\n{len(have)} model(s) already present, {len(to_download)} to download ({human(total)}), "
              f"{len(not_found)} not in repo; {len(node_installs)} node pack(s) to install.")
    else:
        total = 0
        print(f"\nNo model references found; {len(node_installs)} node pack(s) to install.")

    if not to_download and not node_installs:
        print(f"{C_GREEN}Nothing to do — all set!{C_RESET}")
        return
    if dry_run:
        print("(dry run — nothing installed or downloaded)")
        return

    if sys.stdin.isatty() and not assume_yes:
        parts = []
        if node_installs:
            parts.append(f"install {len(node_installs)} node pack(s)")
        if to_download:
            parts.append(f"download {len(to_download)} file(s) ({human(total)})")
        resp = input(f"\nProceed: {' and '.join(parts)}? [Y/n] ").strip().lower()
        if resp not in ("", "y", "yes"):
            sys.exit("Aborted.")

    # -------- install custom-node packs --------
    pack_failures = []
    for name, url, ver, _source in node_installs:
        try:
            install_pack(name, url, ver)
        except Exception as e:
            pack_failures.append((name, str(e)))
            print(f"  {C_RED}✘ failed to install {name}: {e}{C_RESET}")

    # -------- download models --------
    if to_download:
        # hf_transfer parallelizes chunks within a file; snapshot_download parallelizes across files.
        print(f"\n{C_CYAN}Downloading with hf_transfer (parallel)...{C_RESET}\n")
        from huggingface_hub import snapshot_download
        patterns = [m for _, m, _, _ in to_download]
        heartbeat = download_heartbeat(to_download)
        try:
            snapshot_download(
                repo_id=REPO_ID,
                allow_patterns=patterns,
                local_dir=COMFY_DIR,
                token=os.environ.get("HF_TOKEN"),
                max_workers=8,
            )
        finally:
            heartbeat.set()

    print(f"\n{C_BOLD}Verifying:{C_RESET}")
    ok = True
    for name, url, ver, _source in node_installs:
        dest = os.path.join(CUSTOM_NODES_DIR, url.rstrip("/").split("/")[-1].removesuffix(".git"))
        if os.path.isdir(dest) and not any(name == f[0] for f in pack_failures):
            print(f"  {C_GREEN}✔{C_RESET} custom_nodes/{os.path.basename(dest)}")
        else:
            ok = False
            print(f"  {C_RED}✘{C_RESET} {name}  (install failed)")
    for _, match, size, dest in to_download:
        if os.path.exists(dest) and (size == 0 or os.path.getsize(dest) == size):
            print(f"  {C_GREEN}✔{C_RESET} {dest}  ({human(os.path.getsize(dest))})")
        else:
            ok = False
            print(f"  {C_RED}✘{C_RESET} {dest}  (incomplete or missing)")
    if not_found:
        print(f"\n{C_YELLOW}Note:{C_RESET} {len(not_found)} model(s) were not found in the repo (listed above) — "
              f"you'll need to source those elsewhere.")
    if node_installs and ok:
        print(f"\n{C_CYAN}Restart ComfyUI to load the newly installed node packs.{C_RESET}")
    print(f"\n{C_GREEN if ok else C_RED}{'Done — all installs and downloads verified.' if ok else 'Done, but some steps failed — re-run to retry.'}{C_RESET}")

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        sys.exit("\nInterrupted.")
PYEOF
)
# -u: unbuffered stdout, so progress lines appear immediately even when piped (e.g. | tee)
exec python3 -u -c "$PYCODE" "$@"