#!/bin/bash
# Script de téléchargement des meilleurs modèles SDXL
# Pour ComfyUI avec RX 7800 XT
# Version ultra-simple

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
MODELS_DIR="$SCRIPT_DIR/models/checkpoints"
VAE_DIR="$SCRIPT_DIR/models/vae"

mkdir -p "$MODELS_DIR" "$VAE_DIR"

# Fonction simple
download_model() {
    local url=$1
    local output=$2
    local name=$3
    local size=$4
    
    if [ -f "$output" ]; then
        local filesize=$(du -h "$output" | cut -f1)
        echo "[OK] $name deja present ($filesize)"
        return 0
    fi
    
    echo ""
    echo "----------------------------------------"
    echo "Telechargement: $name ($size)"
    echo "Fichier: $(basename "$output")"
    echo "----------------------------------------"
    echo "Patientez... (ne quittez pas)"
    echo ""
    
    # Telechargement avec curl (plus propre que wget)
    if curl -L --progress-bar "$url" -o "$output" 2>&1; then
        echo ""
        echo "[OK] $name telecharge avec succes"
        return 0
    else
        echo ""
        echo "[ERREUR] Probleme lors du telechargement de $name"
        rm -f "$output"
        return 1
    fi
}

# Verif permissions
if [ ! -w "$MODELS_DIR" ]; then
    echo "[ERREUR] Pas de permission d'ecriture dans $MODELS_DIR"
    echo "Execute: sudo chown -R \$USER:\$USER $MODELS_DIR"
    exit 1
fi

# Menu
clear
echo "=========================================="
echo "  GESTIONNAIRE DE MODELES SDXL"
echo "=========================================="
echo ""
echo "Dossier: $MODELS_DIR"
echo "Espace dispo: $(df -h "$MODELS_DIR" | tail -1 | awk '{print $4}')"
echo ""
echo "MODELES DISPONIBLES:"
echo ""
echo "  1) RealVisXL V4.0       - Portraits realistes"
echo "  2) Juggernaut XL v8     - Polyvalent"
echo "  3) DreamShaper XL 1.0   - Art/Fantasy"
echo "  4) SDXL Lightning       - Rapide (8 etapes)"
echo "  5) SDXL Base 1.0        - Modele de base"
echo "  6) SDXL VAE             - Meilleures couleurs"
echo ""
echo "  8) TOUT TELECHARGER     - ~28GB total"
echo "  0) Quitter"
echo ""
echo "=========================================="
read -p "Votre choix (0-8): " choice

case $choice in
    1)
        download_model \
            "https://huggingface.co/SG161222/RealVisXL_V4.0/resolve/main/RealVisXL_V4.0.safetensors" \
            "$MODELS_DIR/RealVisXL_V4.0.safetensors" \
            "RealVisXL V4.0" \
            "6.9GB"
        ;;
    2)
        download_model \
            "https://huggingface.co/RunDiffusion/Juggernaut-XL-v8/resolve/main/juggernautXL_v8Rundiffusion.safetensors" \
            "$MODELS_DIR/juggernautXL_v8Rundiffusion.safetensors" \
            "Juggernaut XL v8" \
            "6.9GB"
        ;;
    3)
        download_model \
            "https://huggingface.co/Lykon/DreamShaper/resolve/main/DreamShaperXL1.0.safetensors" \
            "$MODELS_DIR/DreamShaperXL1.0.safetensors" \
            "DreamShaper XL 1.0" \
            "6.9GB"
        ;;
    4)
        download_model \
            "https://huggingface.co/ByteDance/SDXL-Lightning/resolve/main/sdxl_lightning_8step.safetensors" \
            "$MODELS_DIR/sdxl_lightning_8step.safetensors" \
            "SDXL Lightning 8-step" \
            "6.9GB"
        ;;
    5)
        download_model \
            "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors" \
            "$MODELS_DIR/sd_xl_base_1.0.safetensors" \
            "SDXL Base 1.0" \
            "6.9GB"
        ;;
    6)
        download_model \
            "https://huggingface.co/stabilityai/sdxl-vae/resolve/main/diffusion_pytorch_model.safetensors" \
            "$VAE_DIR/sdxl_vae.safetensors" \
            "SDXL VAE" \
            "335MB"
        ;;
    8)
        echo ""
        echo "=========================================="
        echo "  TELECHARGEMENT DE TOUS LES MODELES"
        echo "  (~28GB au total)"
        echo "=========================================="
        read -p "Confirmer (tapez 'oui'): " confirm
        if [ "$confirm" = "oui" ]; then
            download_model "https://huggingface.co/SG161222/RealVisXL_V4.0/resolve/main/RealVisXL_V4.0.safetensors" "$MODELS_DIR/RealVisXL_V4.0.safetensors" "RealVisXL V4.0" "6.9GB"
            download_model "https://huggingface.co/RunDiffusion/Juggernaut-XL-v8/resolve/main/juggernautXL_v8Rundiffusion.safetensors" "$MODELS_DIR/juggernautXL_v8Rundiffusion.safetensors" "Juggernaut XL v8" "6.9GB"
            download_model "https://huggingface.co/Lykon/DreamShaper/resolve/main/DreamShaperXL1.0.safetensors" "$MODELS_DIR/DreamShaperXL1.0.safetensors" "DreamShaper XL 1.0" "6.9GB"
            download_model "https://huggingface.co/ByteDance/SDXL-Lightning/resolve/main/sdxl_lightning_8step.safetensors" "$MODELS_DIR/sdxl_lightning_8step.safetensors" "SDXL Lightning 8-step" "6.9GB"
            download_model "https://huggingface.co/stabilityai/sdxl-vae/resolve/main/diffusion_pytorch_model.safetensors" "$VAE_DIR/sdxl_vae.safetensors" "SDXL VAE" "335MB"
            echo ""
            echo "=========================================="
            echo "  [OK] TOUS LES MODELES TELECHARGES"
            echo "=========================================="
        else
            echo "Annulation..."
            exit 0
        fi
        ;;
    0)
        echo "Au revoir !"
        exit 0
        ;;
    *)
        echo "[ERREUR] Choix invalide"
        exit 1
        ;;
esac

# Resume
echo ""
echo "=========================================="
echo "  RESUME"
echo "=========================================="
echo ""
echo "Modeles disponibles:"
ls -lh "$MODELS_DIR"/*.safetensors 2>/dev/null | awk '{print "  " $9 " (" $5 ")"}' || echo "  Aucun"
echo ""
echo "VAE disponibles:"
ls -lh "$VAE_DIR"/*.safetensors 2>/dev/null | awk '{print "  " $9 " (" $5 ")"}' || echo "  Aucun"
echo ""
echo "=========================================="
echo "  Prochaine etape:"
echo "  sudo docker restart comfyui"
echo "=========================================="