#!/bin/bash
# Ensure directories exist (relative to script location, e.g., repo root)
mkdir -p ./models/checkpoints
mkdir -p ./models/controlnet
mkdir -p ./models/embeddings
mkdir -p ./models/loras
mkdir -p ./models/motion
mkdir -p ./models/text_encoders
mkdir -p ./models/upscalers
mkdir -p ./output
mkdir -p ./workflows

# Install Hugging Face CLI if not present
if ! command -v huggingface-cli &> /dev/null; then
    pip install huggingface_hub
fi

# Set Hugging Face API token (via environment variable or default for testing)
export HF_TOKEN=${HF_TOKEN:-"your_huggingface_api_token"}  # Replace or set in SimplePod env vars

# Download all models from your Hugging Face repo
echo "Downloading models from your Hugging Face repository..."
huggingface-cli download --token ${HF_TOKEN} myusername/comfyui-dance-models models/checkpoints/sd3.5-base.safetensors --local-dir ./models/checkpoints
huggingface-cli download --token ${HF_TOKEN} myusername/comfyui-dance-models models/loras/dance_lora.safetensors --local-dir ./models/loras
huggingface-cli download --token ${HF_TOKEN} myusername/comfyui-dance-models models/motion/mm_sd_v15_v2.ckpt --local-dir ./models/motion
huggingface-cli download --token ${HF_TOKEN} myusername/comfyui-dance-models models/vae/vae-ft-mse-840000-ema-pruned.safetensors --local-dir ./models/vae
huggingface-cli download --token ${HF_TOKEN} myusername/comfyui-dance-models models/controlnet/control_v11p_sd15_openpose.pth --local-dir ./models/controlnet
huggingface-cli download --token ${HF_TOKEN} myusername/comfyui-dance-models models/controlnet/control_v11p_sd15_canny.pth --local-dir ./models/controlnet

# Install ComfyUI custom nodes for SD 3.5 and motion (if needed)
echo "Installing ComfyUI custom nodes for SD 3.5 and motion..."
pip install git+https://github.com/civitai/sd3-comfyui.git
git clone https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite.git ./models/ComfyUI/custom_nodes/ComfyUI-VideoHelperSuite

# Ensure Rclone is installed (likely pre-installed in ai-dock/comfyui)
if ! command -v rclone &> /dev/null; then
    curl https://rclone.org/install.sh | sudo bash
fi

# Sync outputs from R2 on startup (optional, can be run later)
rclone copy r2bucket:output ./output --create-empty-src-dirs

# Start background sync loop for outputs to R2
nohup bash -c "while true; do rclone sync ./output r2bucket:output; sleep 3600; done" &

# Note: ComfyUI launch is not run locally; this is for SimplePod deployment
echo "Models and structure ready for ComfyUI deployment. Push to Hugging Face and use in SimplePod."