#!/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, with progress and automatic stall recovery.
#
# 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
#   ./download_missing_models.sh --last          # reuse the workflow selection from the previous run
#   ./download_missing_models.sh --parallel true # download up to 8 files at once (also accepts a
#                                                # count, e.g. --parallel 4; default: one at a time)
#
# Auth: uses the HUGGING_FACE_ACCESS_TOKEN environment variable if set
# (falls back to HF_TOKEN), otherwise prompts for a token.
# If GITHUB_TOKEN (or GH_TOKEN) is set, git clones of github.com repos are
# authenticated with it — datacenter IPs (e.g. RunPod) often get 403s on
# anonymous github.com traffic. Use a fine-grained PAT with no repo
# permissions; the token is sent as a one-shot header, never written to disk.
# With or without a token, a failed clone falls back to a repo tarball from
# codeload.github.com, which is typically not IP-blocked.

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# The pod image exports HF_HUB_ENABLE_HF_TRANSFER=1 and HF_XET_HIGH_PERFORMANCE=1.
# hf_transfer is deprecated in huggingface_hub 1.x, and the Xet high-performance
# mode has deadlocked mid-download on this pod (~115 parallel connections, all
# data buffered in RAM, zero bytes written, no error). Stock Xet concurrency is
# still several hundred MB/s here, which is what the network volume tops out at.
unset HF_HUB_ENABLE_HF_TRANSFER HF_XET_HIGH_PERFORMANCE
# 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 downloads Xet-backed repos with hf_xet;
# without it the hub falls back to slower, less robust code paths.
if ! python3 -c "import huggingface_hub" >/dev/null 2>&1; then
    echo "Installing huggingface_hub..."
    pip install -q -U huggingface_hub
fi
python3 -c "import hf_xet" >/dev/null 2>&1 || pip install -q hf_xet >/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")
LAST_FILE = os.path.join(WORKSPACE, ".cache", "last-workflows.json")
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

# ---------------- implicit models (IPAdapterUnifiedLoader) ----------------
#
# IPAdapterUnifiedLoader never names its model files in the workflow JSON: at
# runtime it scans models/clip_vision and models/ipadapter for filenames
# hardcoded per preset (see ComfyUI_IPAdapter_plus/utils.py) and raises
# "ClipVision model not found" if they're absent — so the widget scan above
# can't see them. Map each preset to the files the loader will look for.
# Both SD1.5 and SDXL variants are listed where they exist, because the
# checkpoint architecture is only known at runtime; the unused variant costs
# little (SD1.5 ipadapters are ~100 MB).
CLIP_VIT_H = "clip_vision/CLIP-ViT-H-14-laion2B-s32B-b79K.safetensors"
CLIP_VIT_BIGG = "clip_vision/CLIP-ViT-bigG-14-laion2B-39B-b160k.safetensors"
UNIFIED_LOADER_PRESETS = {
    "LIGHT - SD1.5 only (low strength)": [CLIP_VIT_H, "ipadapter/ip-adapter_sd15_light_v11.bin"],
    "STANDARD (medium strength)": [CLIP_VIT_H, "ipadapter/ip-adapter_sd15.safetensors",
                                   "ipadapter/ip-adapter_sdxl_vit-h.safetensors"],
    "VIT-G (medium strength)": [CLIP_VIT_BIGG, "ipadapter/ip-adapter_sd15_vit-G.safetensors",
                                "ipadapter/ip-adapter_sdxl.safetensors"],
    "PLUS (high strength)": [CLIP_VIT_H, "ipadapter/ip-adapter-plus_sd15.safetensors",
                             "ipadapter/ip-adapter-plus_sdxl_vit-h.safetensors"],
    "PLUS FACE (portraits)": [CLIP_VIT_H, "ipadapter/ip-adapter-plus-face_sd15.safetensors",
                              "ipadapter/ip-adapter-plus-face_sdxl_vit-h.safetensors"],
    "FULL FACE - SD1.5 only (portraits stronger)": [CLIP_VIT_H, "ipadapter/ip-adapter-full-face_sd15.safetensors"],
}

def extract_implicit_models(data):
    """Refs for models that nodes locate on disk at runtime without naming
    them in the JSON. Currently covers IPAdapterUnifiedLoader; the FaceID
    unified loader (insightface + lora files) is not handled."""
    refs = set()
    for n in iter_graph_nodes(data):
        if (n.get("type") or n.get("class_type")) != "IPAdapterUnifiedLoader":
            continue
        inputs = n.get("inputs")
        if isinstance(inputs, dict) and isinstance(inputs.get("preset"), str):
            preset = inputs["preset"]  # API format
        else:  # UI format: the preset is one of the widget values
            wv = n.get("widgets_values") or []
            preset = next((v for v in wv if isinstance(v, str) and v in UNIFIED_LOADER_PRESETS), None)
        if preset in UNIFIED_LOADER_PRESETS:
            refs.update(UNIFIED_LOADER_PRESETS[preset])
        else:
            print(f"  {C_YELLOW}⚠ IPAdapterUnifiedLoader with unrecognized preset {preset!r} — "
                  f"cannot plan its clip_vision/ipadapter files{C_RESET}")
    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"}
GITHUB_TOKEN = os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN")

def github_owner_repo(url):
    m = re.match(r"https?://github\.com/([^/]+)/([^/]+?)(?:\.git)?/?$", url)
    return (m.group(1), m.group(2)) if m else None

def gh_headers():
    h = {"User-Agent": "comfy-dep-installer"}
    if GITHUB_TOKEN:
        h["Authorization"] = f"Bearer {GITHUB_TOKEN}"
    return h

def git_clone(url, dest):
    """Clone, authenticating github.com with GITHUB_TOKEN when available.
    The token goes in a per-invocation header (git -c), so it is never
    persisted in the clone's .git/config."""
    cmd = ["git"]
    if GITHUB_TOKEN and github_owner_repo(url):
        import base64
        basic = base64.b64encode(f"x-access-token:{GITHUB_TOKEN}".encode()).decode()
        cmd += ["-c", f"http.https://github.com/.extraheader=Authorization: basic {basic}"]
    cmd += ["clone", "--recursive", "--progress", url, dest]
    subprocess.run(cmd, check=True, env=GIT_ENV, timeout=900, stdin=subprocess.DEVNULL)

def install_from_tarball(url, ver, dest):
    """Fallback when git clone is refused (github.com 403s anonymous requests
    from many datacenter IPs): fetch the repo tarball from codeload.github.com,
    which sits on separate infrastructure and is typically not blocked.
    Caveats vs a real clone: no .git (ComfyUI-Manager can't update the pack)
    and no submodules."""
    owner, repo = github_owner_repo(url)
    if ver and re.fullmatch(r"[0-9a-f]{40}", str(ver), re.I):
        ref = str(ver)
    else:
        info = fetch_json_gh(f"https://api.github.com/repos/{owner}/{repo}")
        ref = info.get("default_branch") or "main"
    print(f"  {C_YELLOW}falling back to tarball of {owner}/{repo}@{ref[:12]} via codeload.github.com ...{C_RESET}")
    import shutil, tarfile, tempfile
    req = urllib.request.Request(
        f"https://codeload.github.com/{owner}/{repo}/tar.gz/{ref}", headers=gh_headers())
    os.makedirs(CUSTOM_NODES_DIR, exist_ok=True)
    # tmp dir lives next to dest so the final rename stays on one filesystem;
    # "." prefix keeps installed_packs() from ever seeing it
    with tempfile.TemporaryDirectory(dir=CUSTOM_NODES_DIR, prefix=".tarball-") as tmp:
        tar_path = os.path.join(tmp, "repo.tar.gz")
        with urllib.request.urlopen(req, timeout=120) as r, open(tar_path, "wb") as fh:
            shutil.copyfileobj(r, fh)
        with tarfile.open(tar_path) as tar:
            try:
                tar.extractall(tmp, filter="data")
            except TypeError:  # Python < 3.12 has no filter=
                tar.extractall(tmp)
        tops = [e for e in os.listdir(tmp)
                if os.path.isdir(os.path.join(tmp, e))]
        if len(tops) != 1:
            raise RuntimeError(f"unexpected tarball layout: {tops}")
        os.rename(os.path.join(tmp, tops[0]), dest)

def fetch_json_gh(url):
    req = urllib.request.Request(url, headers=gh_headers())
    with urllib.request.urlopen(req, timeout=30) as r:
        return json.load(r)

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):
        try:
            git_clone(url, dest)
        except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e:
            if not github_owner_repo(url):
                raise
            print(f"  {C_YELLOW}⚠ git clone failed ({e}){C_RESET}")
            install_from_tarball(url, ver, dest)
    if ver and re.fullmatch(r"[0-9a-f]{40}", str(ver), re.I) \
            and os.path.isdir(os.path.join(dest, ".git")):
        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 parse_parallel(argv):
    """Consume --parallel false|true|<N> (or --parallel=<val>) from argv and
    return (remaining argv, max_workers). Default is one file at a time —
    pass --parallel true (or a count) for concurrent files. Must run before
    pick_workflows, which would otherwise mistake the flag's value for a
    workflow name."""
    out, max_workers, i = [], 1, 0
    while i < len(argv):
        a = argv[i]
        if a == "--parallel" or a.startswith("--parallel="):
            if "=" in a:
                val = a.split("=", 1)[1]
            else:
                i += 1
                val = argv[i] if i < len(argv) else ""
            v = val.strip().lower()
            if v in ("false", "no", "off"):
                max_workers = 1
            elif v in ("true", "yes", "on", ""):
                max_workers = 8
            elif v.isdigit() and int(v) > 0:
                max_workers = int(v)
            else:
                sys.exit(f"--parallel expects true/false or a positive worker count, got '{val}'")
        else:
            out.append(a)
        i += 1
    return out, max_workers

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 bytes_on_disk(to_download):
    """Bytes actually landed for the planned downloads: finished files,
    hub .incomplete fragments, and our curl .part files."""
    done = 0
    cache = os.path.join(COMFY_DIR, ".cache", "huggingface", "download")
    for _, _, size, dest in to_download:
        if os.path.exists(dest):
            done += min(os.path.getsize(dest), size)
        if os.path.exists(dest + ".part"):
            done += os.path.getsize(dest + ".part")
    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
    return done

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()
    def run():
        while not stop.wait(20):
            done = bytes_on_disk(to_download)
            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

STALL_SECS = 10          # once bytes are flowing, a freeze this long = wedged, kill it
STARTUP_GRACE_SECS = 60  # each (re)launch gets this long to land its first bytes (metadata resolution etc.)
RESTART_WAIT_SECS = 2    # pause between killing a wedged download and relaunching it
MAX_RESTARTS = 3         # kill-and-restart attempts before falling back to curl

def hf_download_with_watchdog(patterns, to_download, max_workers):
    """Run snapshot_download in a child process and kill it if bytes-on-disk
    freeze. Both hf_transfer and the Xet backend have wedged mid-file on this
    pod (data buffered in RAM, file frozen, no error raised), so a blocking
    in-process call can hang forever. Any change in bytes-on-disk counts as
    activity (a restarted attempt may truncate a leftover .incomplete
    fragment, briefly shrinking the total). Returns True if the hub download
    finished."""
    import multiprocessing as mp
    def child():
        from huggingface_hub import snapshot_download
        snapshot_download(
            repo_id=REPO_ID,
            allow_patterns=patterns,
            local_dir=COMFY_DIR,
            token=os.environ.get("HF_TOKEN"),
            max_workers=max_workers,
        )
    p = mp.Process(target=child)
    p.start()
    last, last_t, started = bytes_on_disk(to_download), time.time(), False
    while p.is_alive():
        p.join(timeout=2)
        if not p.is_alive():
            break
        done = bytes_on_disk(to_download)
        if done != last:
            last, last_t, started = done, time.time(), True
        elif time.time() - last_t > (STALL_SECS if started else STARTUP_GRACE_SECS):
            print(f"\n  {C_YELLOW}⚠ no bytes written for {int(time.time() - last_t)}s — hub download is wedged, killing it{C_RESET}", flush=True)
            p.terminate()
            p.join(10)
            if p.is_alive():
                p.kill()
                p.join()
            return False
    return p.exitcode == 0

def still_missing(to_download):
    return [t for t in to_download
            if not (os.path.exists(t[3]) and (t[2] == 0 or os.path.getsize(t[3]) == t[2]))]

def download_models(to_download, max_workers):
    """Hub download with automatic stall recovery: a wedged download is
    killed, we wait RESTART_WAIT_SECS, and relaunch the same hub download
    (files that already finished are skipped). After MAX_RESTARTS relaunches,
    whatever is still missing goes through plain resumable HTTP with curl."""
    patterns = [m for _, m, _, _ in to_download]
    heartbeat = download_heartbeat(to_download)
    try:
        for attempt in range(MAX_RESTARTS + 1):
            if attempt:
                left = still_missing(to_download)
                print(f"  {C_CYAN}waiting {RESTART_WAIT_SECS}s, then restarting the hub download "
                      f"({len(left)} file(s) to go) — retry {attempt} of {MAX_RESTARTS} ...{C_RESET}", flush=True)
                time.sleep(RESTART_WAIT_SECS)
            if hf_download_with_watchdog(patterns, to_download, max_workers) or not still_missing(to_download):
                return
    finally:
        heartbeat.set()
    print(f"\n{C_YELLOW}Hub download still incomplete after {MAX_RESTARTS} restart(s) — switching to plain resumable HTTP.{C_RESET}")
    curl_fallback(to_download)

def curl_fallback(to_download):
    """Plain resumable HTTP for whatever is still missing. --speed-limit makes
    curl abort any transfer that drops below 1 MB/s for 30s (instead of
    hanging on a dead connection), and -C - resumes from the .part file."""
    token = os.environ.get("HF_TOKEN", "")
    for _, match, size, dest in to_download:
        if os.path.exists(dest) and (size == 0 or os.path.getsize(dest) == size):
            continue
        url = f"https://huggingface.co/{REPO_ID}/resolve/main/{match}"
        part = dest + ".part"
        os.makedirs(os.path.dirname(dest), exist_ok=True)
        print(f"  {C_CYAN}fetching {match} with curl (resumable) ...{C_RESET}", flush=True)
        for attempt in range(1, 11):
            r = subprocess.run(
                ["curl", "-L", "--fail", "-C", "-",
                 "-H", f"Authorization: Bearer {token}",
                 "--speed-limit", "1000000", "--speed-time", "30",
                 "--progress-bar", "-o", part, url],
                stdin=subprocess.DEVNULL)
            have = os.path.getsize(part) if os.path.exists(part) else 0
            if r.returncode == 0 and (size == 0 or have == size):
                os.replace(part, dest)
                break
            print(f"  {C_YELLOW}… attempt {attempt} stopped at {human(have)} of {human(size)} — resuming{C_RESET}", flush=True)
        else:
            print(f"  {C_RED}✘ gave up on {match} after 10 attempts{C_RESET}")

def main():
    argv, max_workers = parse_parallel(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}")

    if "--last" in argv:
        try:
            with open(LAST_FILE) as fh:
                last_names = json.load(fh)
        except (OSError, ValueError):
            last_names = None
        if not last_names:
            sys.exit("--last: no saved selection found — run once and pick workflows first.")
        by_name = {os.path.basename(f): f for f in files}
        gone = [n for n in last_names if n not in by_name]
        if gone:
            sys.exit(f"--last: saved workflow(s) no longer in {WORKFLOWS_DIR}: {', '.join(gone)}")
        chosen = [by_name[n] for n in last_names]
        print(f"{C_CYAN}Reusing last selection (--last).{C_RESET}")
    else:
        chosen = pick_workflows(files, argv)
        # remember the selection so a retry can just use --last
        try:
            os.makedirs(os.path.dirname(LAST_FILE), exist_ok=True)
            with open(LAST_FILE, "w") as fh:
                json.dump([os.path.basename(f) for f in chosen], fh)
        except OSError:
            pass
    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)
        implicit = extract_implicit_models(data) - refs
        suffix = f" + {len(implicit)} implicit via IPAdapterUnifiedLoader" if implicit else ""
        print(f"  • {os.path.basename(f)}  ({len(refs)} model refs{suffix})")
        for r in refs | implicit:
            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:
        mode = "one file at a time" if max_workers == 1 else f"up to {max_workers} files in parallel"
        print(f"\n{C_CYAN}Downloading with the Hugging Face hub (Xet backend, {mode})...{C_RESET}\n")
        download_models(to_download, max_workers)

    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" "$@"