#!/usr/bin/env bash
# One-command pull for keys-2k MiniMax-H3 Parallel (Two DGX Sparks)
# Usage:
#   curl -fsSL https://huggingface.co/drowzeys/keys-2k-MiniMax-H3-Parallel-Two-DGX-Sparks/raw/main/scripts/one_command_pull.sh | bash
# Or:
#   bash one_command_pull.sh
set -euo pipefail

HF_REPO="${HF_REPO:-drowzeys/keys-2k-MiniMax-H3-Parallel-Two-DGX-Sparks}"
IMAGE="${IMAGE:-ghcr.io/drowzeys/keys-2k-minimax-h3-parallel-two-dgx-sparks:0.31.1-pp20260811}"
INSTALL_DIR="${INSTALL_DIR:-$HOME/keys-2k-minimax-h3-parallel}"
MODELS_DIR="${MODELS_DIR:-$INSTALL_DIR/models}"

echo "=== keys-2k MiniMax-H3 Parallel — one-command pull ==="
echo "  HF:     $HF_REPO"
echo "  image:  $IMAGE"
echo "  install:$INSTALL_DIR"
echo

mkdir -p "$INSTALL_DIR" "$MODELS_DIR"

# 1) Docker image
if ! command -v docker >/dev/null; then
  echo "ERROR: docker not found" >&2
  exit 1
fi
echo "--> docker pull $IMAGE"
if ! docker pull "$IMAGE"; then
  echo "WARN: ghcr pull failed — try: echo \$GHCR_TOKEN | docker login ghcr.io -u USER --password-stdin" >&2
  echo "      or load from HF release asset if published" >&2
fi
docker tag "$IMAGE" keys-2k-minimax-h3-parallel:latest 2>/dev/null || true

# 2) Package files from HF (compose, workflows, scripts)
if command -v hf >/dev/null; then
  echo "--> hf download $HF_REPO → $INSTALL_DIR"
  hf download "$HF_REPO" --local-dir "$INSTALL_DIR" --exclude "models/**" 2>/dev/null \
    || huggingface-cli download "$HF_REPO" --local-dir "$INSTALL_DIR" 2>/dev/null \
    || true
elif command -v git >/dev/null; then
  if [[ ! -d "$INSTALL_DIR/.git" ]]; then
    git clone "https://huggingface.co/$HF_REPO" "$INSTALL_DIR" 2>/dev/null || true
  fi
fi

chmod +x "$INSTALL_DIR"/scripts/*.sh 2>/dev/null || true

# 3) Weight skeleton + fetch helper
mkdir -p "$MODELS_DIR"/{diffusion_models,text_encoders/H3,vae,upscale_models,loras}
if [[ -x "$INSTALL_DIR/scripts/fetch_weights.sh" ]]; then
  echo "--> optional weights (set FETCH_WEIGHTS=1 to download large files now)"
  if [[ "${FETCH_WEIGHTS:-0}" == "1" ]]; then
    MODELS_DIR="$MODELS_DIR" bash "$INSTALL_DIR/scripts/fetch_weights.sh"
  else
    echo "    skip weights. Later: FETCH_WEIGHTS=1 bash $INSTALL_DIR/scripts/fetch_weights.sh"
    echo "    or bind-mount existing models from your lab tree."
  fi
fi

cat <<EOF

=== ready ===
Single node (after models present under $MODELS_DIR):

  docker run --rm --gpus all --network host \\
    -v "$MODELS_DIR:/opt/powerpack-h3/ComfyUI/models" \\
    -v "\$HOME/Videos/h3-out:/opt/powerpack-h3/ComfyUI/output" \\
    $IMAGE

Dual Sparks (parallel 2K bench):

  export HEAD=10.100.10.1 WORKER=10.100.10.5
  export MODELS_DIR=$MODELS_DIR
  bash $INSTALL_DIR/scripts/run_dual_h3.sh
  bash $INSTALL_DIR/scripts/run_anime_2k_bench.sh

Manifest: $INSTALL_DIR/manifest/H3_STACK_MANIFEST.json
Docs:     https://huggingface.co/$HF_REPO
EOF