#!/bin/bash
# Pull the freshest workflows/, download_missing_models.sh, and README.md from
# the HF repo (aleph65/ComfyUI) into /workspace on every boot, so pods never
# run against the stale copies baked into the image or left on a network
# volume. Best-effort: on any failure (offline, bad token, read-only
# /workspace) the existing copies stay in place and startup never blocks.
# Set WORKSPACE_REFRESH=false in the pod env to skip.

REPO_ID="aleph65/ComfyUI"
TAG="[refresh-workspace]"
TOKEN="${HUGGING_FACE_ACCESS_TOKEN:-${HF_TOKEN:-}}"

if [[ "${WORKSPACE_REFRESH:-true}" == "false" ]]; then
    echo "$TAG disabled via WORKSPACE_REFRESH=false — skipping."
    exit 0
fi
if [[ ! -d /workspace || ! -w /workspace ]]; then
    echo "$TAG /workspace missing or not writable — skipping." >&2
    exit 0
fi

set -u

# 1) single files: atomic swap, keep the old copy on any failure
update_workspace_file() {
    local name="$1" check="$2"
    local url="https://huggingface.co/$REPO_ID/resolve/main/$name"
    local dest="/workspace/$name"
    local tmp auth=()
    tmp=$(mktemp) || return 0
    [[ -n "$TOKEN" ]] && auth=(-H "Authorization: Bearer $TOKEN")
    if curl -fsSL --connect-timeout 10 --max-time 60 "${auth[@]}" "$url" -o "$tmp" \
        && head -c 2 "$tmp" | grep -q "$check"; then
        if ! cmp -s "$tmp" "$dest" 2>/dev/null; then
            # scripts (shebang in first bytes) must stay executable, even
            # without a .sh suffix (e.g. start_comfy); everything else 644
            if head -c2 "$tmp" | grep -q '#!'; then chmod 755 "$tmp"; else chmod 644 "$tmp"; fi
            mv "$tmp" "$dest"
            echo "$TAG updated $name from hf.co/$REPO_ID."
        fi
    else
        echo "$TAG could not fetch latest $name — keeping existing copy." >&2
    fi
    rm -f "$tmp"
}
update_workspace_file download_missing_models.sh '#!' || true
update_workspace_file README.md '.' || true
update_workspace_file start_comfy '#!' || true
update_workspace_file claude_start.sh '#!' || true

# 2) mirror workflows/ from the repo (temp snapshot + rsync --delete, so
# renamed/removed workflows disappear too; local-only edits do NOT survive —
# the repo is the source of truth)
TMP_DIR=$(mktemp -d)
if timeout 240 python3 - "$REPO_ID" "$TMP_DIR" "$TOKEN" <<'PY'
import sys
repo_id, tmp, token = sys.argv[1], sys.argv[2], sys.argv[3]
from huggingface_hub import snapshot_download
snapshot_download(repo_id=repo_id, allow_patterns=["workflows/*"],
                  local_dir=tmp, token=token or None)
PY
then
    if [[ -d "$TMP_DIR/workflows" ]]; then
        rsync -a --delete --exclude='.cache' "$TMP_DIR/workflows/" /workspace/workflows/
        echo "$TAG workflows/ synced from hf.co/$REPO_ID ($(find /workspace/workflows -name '*.json' | wc -l) json files)."
    else
        echo "$TAG snapshot returned no workflows/ — keeping existing copies." >&2
    fi
else
    echo "$TAG workflows sync failed — keeping existing copies." >&2
fi
rm -rf "$TMP_DIR"