#!/bin/bash
# ==============================================================================
# Oz Videos — Vast.ai Provisioning Script
# All assets from parlorsky/instara-models (v3-models/)
# ==============================================================================
set +e

HF_REPO="parlorsky/instara-models"
HF_TOKEN="${HF_TOKEN:-}"
HF_BASE="https://huggingface.co/$HF_REPO/resolve/main"

echo ""
echo "=========================================="
echo "  Oz Videos — Deploy"
echo "=========================================="
echo ""

# ── Phase 0: Locate ComfyUI ──
echo "=== Phase 0: Locating ComfyUI ==="
COMFY_DIR=""
for i in $(seq 1 60); do
    for try_dir in "/opt/workspace-internal/ComfyUI" "/workspace/ComfyUI" "/root/ComfyUI" "/workspace/runpod-slim/ComfyUI"; do
        [ -f "$try_dir/main.py" ] && { COMFY_DIR="$try_dir"; break 2; }
    done
    echo "  Waiting for ComfyUI... ($i)"
    sleep 5
done
[ -z "$COMFY_DIR" ] && COMFY_DIR=$(find /opt /workspace /root -maxdepth 5 -name "main.py" -path "*/ComfyUI/*" 2>/dev/null | head -1 | xargs dirname 2>/dev/null)
if [ -z "$COMFY_DIR" ] || [ ! -f "$COMFY_DIR/main.py" ]; then echo "[!] ComfyUI not found!"; exit 1; fi
echo "[ok] ComfyUI: $COMFY_DIR"

NODES="$COMFY_DIR/custom_nodes"
MODELS="$COMFY_DIR/models"

for p in "/venv/main/bin" "$COMFY_DIR/.venv-cu128/bin" "$COMFY_DIR/.venv/bin"; do
    [ -x "$p/pip" ] && { PIP="$p/pip"; PYTHON="$p/python"; break; }
done
PIP="${PIP:-pip}"; PYTHON="${PYTHON:-python3}"

while pgrep -f "pip install.*requirements" > /dev/null 2>&1; do sleep 5; done
echo "[ok] Template setup complete"

# ── Phase 1: System deps ──
echo ""
echo "=== Phase 1: System dependencies ==="
apt-get update -qq && apt-get install -y -qq aria2 2>/dev/null || true

# ── Phase 2: Custom nodes ──
echo ""
echo "=== Phase 2: Custom nodes ==="
cd "$NODES"

install_node() {
    local name="$1" url="$2"
    [ -d "$name" ] && { echo "  [ok] $name"; return 0; }
    echo "  [+] $name..."
    git clone --depth 1 "$url" "$name" 2>/dev/null || { echo "  [!] $name failed"; return 0; }
    [ -f "$name/requirements.txt" ] && $PIP install -r "$name/requirements.txt" --quiet --no-deps 2>/dev/null || true
    [ -f "$name/install.py" ] && (cd "$name" && $PYTHON install.py 2>/dev/null) || true
}

# V3 PRO nodes
install_node "ComfyUI-WanVideoWrapper" "https://github.com/kijai/ComfyUI-WanVideoWrapper"
install_node "ComfyUI-KJNodes" "https://github.com/kijai/ComfyUI-KJNodes"
install_node "ComfyUI-VideoHelperSuite" "https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite"
install_node "ComfyUI-WanAnimatePreprocess" "https://github.com/kijai/ComfyUI-WanAnimatePreprocess.git"
install_node "rgthree-comfy" "https://github.com/rgthree/rgthree-comfy"
install_node "ComfyUI-Manager" "https://github.com/ltdrdata/ComfyUI-Manager"
install_node "ComfyUI-mxToolkit" "https://github.com/Smirnov75/ComfyUI-mxToolkit"
install_node "ComfyUI-segment-anything-2" "https://github.com/kijai/ComfyUI-segment-anything-2"
install_node "cg-use-everywhere" "https://github.com/chrisgoringe/cg-use-everywhere"
# Teskor Utils (from our HF backup — has TSColorMatch, TSPreviewImageNoMetadata, TSVideoCombineNoMetadata)
if [ ! -d "comfyui-teskors-utils" ]; then
    echo "  [+] comfyui-teskors-utils (from HF)..."
    $PIP install -q huggingface_hub 2>/dev/null
    $PYTHON -c "
from huggingface_hub import snapshot_download
import shutil, os
snapshot_download(repo_id='$HF_REPO', allow_patterns='v3-nodes/comfyui-teskors-utils/*', local_dir='/tmp/ts_dl', token='$HF_TOKEN')
src = '/tmp/ts_dl/v3-nodes/comfyui-teskors-utils'
if os.path.exists(src): shutil.copytree(src, '$NODES/comfyui-teskors-utils')
" 2>/dev/null
    [ -f "$NODES/comfyui-teskors-utils/requirements.txt" ] && $PIP install -q -r "$NODES/comfyui-teskors-utils/requirements.txt" --no-deps 2>/dev/null
    rm -rf /tmp/ts_dl
else
    echo "  [ok] comfyui-teskors-utils"
fi

# Remove any leftover proxy/decoy nodes that shadow real implementations
rm -rf "$NODES/ts-proxy-nodes" "$NODES/NEW-UTILS" "$NODES/comfyui-teskors-utils-main" 2>/dev/null

# KJNodes fix
KJFILE="$NODES/ComfyUI-KJNodes/nodes/nodes.py"
[ -f "$KJFILE" ] && grep -q "search_aliases" "$KJFILE" 2>/dev/null && sed -i 's/search_aliases=\[.*\],\?//g' "$KJFILE"

# Python deps
$PIP install --quiet transformers accelerate "antlr4-python3-runtime==4.9.3" pyparsing matplotlib || true
$PIP install --quiet onnxruntime-gpu 2>/dev/null || $PIP install --quiet onnxruntime || true
echo "[ok] Custom nodes ready"

# ── Phase 3: Download models ──
echo ""
echo "=== Phase 3: Download models (from parlorsky/instara-models) ==="

dl() {
    local dir="$1" file="$2" url="$3"
    mkdir -p "$dir"
    [ -f "$dir/$file" ] && { echo "  [ok] $file"; return; }
    echo "  [dl] $file..."
    local auth=()
    [[ "$url" =~ huggingface\.co ]] && [ -n "$HF_TOKEN" ] && auth=("--header=Authorization: Bearer $HF_TOKEN")
    if command -v aria2c >/dev/null 2>&1; then
        aria2c --console-log-level=error -c -x 16 -s 16 -k 1M "${auth[@]}" -d "$dir" -o "$file" "$url" 2>/dev/null && return
    fi
    curl -s -L -H "Authorization: Bearer $HF_TOKEN" -o "$dir/$file" "$url"
}

# Main diffusion model (17.3GB)
dl "$MODELS/diffusion_models" "ryusei_prime.safetensors" "$HF_BASE/v3-models/WanModel.safetensors"

# CLIP
dl "$MODELS/clip" "sakura_kaze_enc.safetensors" "$HF_BASE/v3-models/text_enc.safetensors"
dl "$MODELS/clip_vision" "hanabi_vis_core.safetensors" "$HF_BASE/v3-models/klip_vision.safetensors"

# VAE
dl "$MODELS/vae" "tsuki_vae_gen.safetensors" "$HF_BASE/v3-models/vae.safetensors"

# LoRAs
dl "$MODELS/loras" "nami_fun_hps21.safetensors" "$HF_BASE/v3-models/WanFun.reworked.safetensors"
dl "$MODELS/loras" "hikari_cfg4_dist.safetensors" "$HF_BASE/v3-models/light.safetensors"
dl "$MODELS/loras" "fuji_pusa_r512.safetensors" "$HF_BASE/v3-models/WanPusa.safetensors"
dl "$MODELS/loras" "kaminari_i2v_rev.safetensors" "$HF_BASE/v3-models/wan.reworked.safetensors"

# Detection / Pose
dl "$MODELS/detection" "yolov10m.onnx" "$HF_BASE/v3-models/detection/yolov10m.onnx"
dl "$MODELS/detection" "vitpose_h_wholebody_model.onnx" "$HF_BASE/v3-models/detection/vitpose_h_wholebody_model.onnx"
dl "$MODELS/detection" "vitpose_h_wholebody_data.bin" "$HF_BASE/v3-models/detection/vitpose_h_wholebody_data.bin"

# Uni3C ControlNet (V3 PRO video control)
dl "$MODELS/controlnet" "tenshi_ctrl_v1.safetensors" "$HF_BASE/v3-models/Wan21_Uni3C_controlnet_fp16.safetensors"

# SAM2 (auto-segmentation)
dl "$MODELS/sam2" "sam2.1_hiera_base_plus.safetensors" "$HF_BASE/v3-models/sam2.1_hiera_base_plus.safetensors"

echo "[ok] All models ready"

# ── Phase 4: Deploy V3 PRO workflow ──
echo ""
echo "=== Phase 4: Deploy V3 PRO workflow ==="
WF_DIR="$COMFY_DIR/user/default/workflows"
mkdir -p "$WF_DIR"

$PIP install -q huggingface_hub 2>/dev/null
$PYTHON -c "
from huggingface_hub import hf_hub_download
import shutil
local = hf_hub_download('$HF_REPO', 'v3/Oz_engine_v3_pro.json', token='$HF_TOKEN')
shutil.copy(local, '$WF_DIR/Oz_V3_PRO.json')
shutil.copy(local, '$COMFY_DIR/user/default/workflows/Oz_V3_PRO.json')
print('[ok] V3 PRO workflow deployed')
" 2>/dev/null || echo "[warn] Workflow download had issues"

# ── Phase 5: Settings ──
echo ""
echo "=== Phase 5: Settings ==="
SETTINGS_FILE="$COMFY_DIR/user/default/comfy.settings.json"
mkdir -p "$(dirname "$SETTINGS_FILE")"
cat > "$SETTINGS_FILE" << 'EOF'
{"Comfy.Locale":"en","Comfy.Workflow.WorkflowTabsPosition":"Sidebar","Comfy.Sidebar.Location":"left","Comfy.Sidebar.Size":"small"}
EOF
echo "[ok] Settings configured (locale: en)"

# ── Phase 6: Restart ComfyUI (must restart even if already running to pick up new nodes) ──
echo ""
echo "=== Phase 6: Restart ComfyUI ==="
# Write a detached restart script so it survives parent shell exit
cat > /tmp/oz_restart.sh << 'RESTARTEOF'
#!/bin/bash
sleep 1
kill -9 $(pgrep -f "main.py") 2>/dev/null
sleep 3
cd COMFY_DIR_PLACEHOLDER
PYTHON_PLACEHOLDER main.py --listen 0.0.0.0 --port 8188 --cuda-device 0 > /tmp/comfyui.log 2>&1
RESTARTEOF
sed -i "s|COMFY_DIR_PLACEHOLDER|$COMFY_DIR|g" /tmp/oz_restart.sh
sed -i "s|PYTHON_PLACEHOLDER|$PYTHON|g" /tmp/oz_restart.sh
chmod +x /tmp/oz_restart.sh
setsid bash /tmp/oz_restart.sh &
disown

for i in $(seq 1 60); do
    [ "$(curl -s -o /dev/null -w '%{http_code}' --max-time 3 http://localhost:8188/system_stats 2>/dev/null)" = "200" ] && { echo "[ok] ComfyUI online (${i}s)"; break; }
    sleep 2
done

echo ""
echo "=========================================="
echo "  Oz Videos — Deploy Complete!"
echo "  ComfyUI: http://localhost:8188"
echo "  Workflow: Oz_Videos.json"
echo "  Disk: ~30GB models"
echo "=========================================="