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

set -e

# Détection robuste du répertoire du script
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"

# Couleurs pour l'affichage
GREEN='\033[0;32m'
BLUE='\033[0;34m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
NC='\033[0m' # No Color

# Fonction d'affichage d'en-tête
print_header() {
    echo ""
    echo "========================================"
    echo "$1"
    echo "========================================"
    echo ""
}

# Fonction de téléchargement simple et propre
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 -e "${GREEN}✓${NC} $name déjà présent ($filesize)"
        return 0
    fi
    
    echo -e "${BLUE}📥${NC} Téléchargement de $name ($size)..."
    echo "   Destination: $(basename "$output")"
    echo ""
    
    # Téléchargement avec progression simple
    if wget --show-progress --progress=dot:giga -O "$output" "$url" 2>&1; then
        echo ""
        echo -e "${GREEN}✓${NC} $name téléchargé avec succès"
        return 0
    else
        echo ""
        echo -e "${RED}✗${NC} Erreur lors du téléchargement de $name"
        rm -f "$output"
        return 1
    fi
}

# Affichage d'accueil
clear
print_header "🎨 GESTIONNAIRE DE MODÈLES SDXL"
echo "Répertoire des modèles: $MODELS_DIR"
echo "Espace disque disponible: $(df -h "$MODELS_DIR" | tail -1 | awk '{print $4}')"
echo ""
echo "Espace requis estimé:"
echo "  • 1 modèle: ~7 GB"
echo "  • 3 modèles: ~21 GB"  
echo "  • Tous les modèles: ~35 GB"
echo ""
read -p "Appuyez sur Entrée pour continuer..."

# Menu principal
clear
print_header "📋 MENU PRINCIPAL"

echo "=== MODÈLES PHOTORÉALISTES ==="
echo ""
echo "1) RealVisXL V4.0 (6.9GB)"
echo "   Portraits ultra-réalistes, peau détaillée"
echo ""
echo "2) Juggernaut XL v9 (6.9GB)"
echo "   Polyvalent, excellent pour tout type d'image"
echo ""
echo "3) Proteus V0.4 (6.9GB)"
echo "   Photoréalisme amélioré, excellente qualité"
echo ""
echo "=== MODÈLES ART/STYLISÉS ==="
echo ""
echo "4) DreamShaper XL 1.0 (6.9GB)"
echo "   Art fantastique, peinture, illustration"
echo ""
echo "5) SDXL Base 1.0 (6.9GB)"
echo "   Modèle de base Stability AI (déjà installé)"
echo ""
echo "=== MODÈLES RAPIDES ==="
echo ""
echo "6) SDXL Lightning 8-step (6.9GB)"
echo "   Génération 4x plus rapide, 8 étapes seulement"
echo ""
echo "=== COMPOSANTS ==="
echo ""
echo "7) SDXL VAE (335MB)"
echo "   Meilleure qualité de couleurs pour tous les modèles"
echo ""
echo "=== INSTALLATION MASSIVE ==="
echo ""
echo "8) TOUT TÉLÉCHARGER (~35GB)"
echo "   Tous les modèles ci-dessus"
echo ""
echo "0) Quitter"
echo ""
echo "========================================"
read -p "Votre choix (0-8): " choice

# Traitement du choix
case $choice in
    1)
        clear
        print_header "📥 TÉLÉCHARGEMENT: RealVisXL V4.0"
        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"
        echo ""
        echo -e "${YELLOW}💡 Conseil d'utilisation:${NC}"
        echo "   • Parfait pour les portraits"
        echo "   • Prompt: 'professional photo, 8k uhd, dslr, soft lighting'"
        echo "   • Étapes: 40-50, CFG: 7-8"
        ;;
    2)
        clear
        print_header "📥 TÉLÉCHARGEMENT: Juggernaut XL v9"
        download_model \
            "https://huggingface.co/RunDiffusion/Juggernaut-XL-v9/resolve/main/Juggernaut-XL-v9.safetensors" \
            "$MODELS_DIR/Juggernaut-XL-v9.safetensors" \
            "Juggernaut XL v9" \
            "6.9GB"
        echo ""
        echo -e "${YELLOW}💡 Conseil d'utilisation:${NC}"
        echo "   • Modèle polyvalent quotidien"
        echo "   • Fonctionne avec tous les styles"
        echo "   • Étapes: 30-40, CFG: 7-8"
        ;;
    3)
        clear
        print_header "📥 TÉLÉCHARGEMENT: Proteus V0.4"
        download_model \
            "https://huggingface.co/dathudehenry/ProteusV0.4/resolve/main/ProteusV0.4.safetensors" \
            "$MODELS_DIR/ProteusV0.4.safetensors" \
            "Proteus V0.4" \
            "6.9GB"
        echo ""
        echo -e "${YELLOW}💡 Conseil d'utilisation:${NC}"
        echo "   • Excellente texture de peau"
        echo "   • Prompt: 'raw photo, realistic, detailed skin texture'"
        echo "   • Étapes: 30-40, CFG: 6-7"
        ;;
    4)
        clear
        print_header "📥 TÉLÉCHARGEMENT: DreamShaper XL 1.0"
        download_model \
            "https://huggingface.co/Lykon/DreamShaper/resolve/main/DreamShaperXL1.0.safetensors" \
            "$MODELS_DIR/DreamShaperXL1.0.safetensors" \
            "DreamShaper XL 1.0" \
            "6.9GB"
        echo ""
        echo -e "${YELLOW}💡 Conseil d'utilisation:${NC}"
        echo "   • Art, fantasy, illustrations"
        echo "   • Prompt: 'digital art, concept art, highly detailed'"
        echo "   • Étapes: 30-40, CFG: 7-8"
        ;;
    5)
        clear
        print_header "📥 TÉLÉCHARGEMENT: SDXL Base 1.0"
        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"
        echo ""
        echo -e "${YELLOW}💡 Conseil d'utilisation:${NC}"
        echo "   • Modèle officiel Stability AI"
        echo "   • Très stable et prévisible"
        echo "   • Étapes: 30-40, CFG: 7-8"
        ;;
    6)
        clear
        print_header "📥 TÉLÉCHARGEMENT: SDXL Lightning"
        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"
        echo ""
        echo -e "${YELLOW}⚠️  IMPORTANT - Réglages spéciaux:${NC}"
        echo "   • Étapes: 8 (OBLIGATOIRE)"
        echo "   • CFG: 1-2 (TRÈS BAS)"
        echo "   • 4x plus rapide que les autres modèles"
        ;;
    7)
        clear
        print_header "📥 TÉLÉCHARGEMENT: SDXL VAE"
        download_model \
            "https://huggingface.co/stabilityai/sdxl-vae/resolve/main/diffusion_pytorch_model.safetensors" \
            "$VAE_DIR/sdxl_vae.safetensors" \
            "SDXL VAE" \
            "335MB"
        echo ""
        echo -e "${YELLOW}💡 Conseil d'utilisation:${NC}"
        echo "   • À utiliser avec tous les modèles SDXL"
        echo "   • Améliore les couleurs et détails"
        echo "   • Chargez-le dans le node 'Load VAE' de ComfyUI"
        ;;
    8)
        clear
        print_header "📥 TÉLÉCHARGEMENT COMPLET"
        echo "Cela va télécharger ~35GB de modèles"
        echo ""
        read -p "Êtes-vous sûr ? (tapez 'oui' pour confirmer): " confirm
        if [ "$confirm" = "oui" ]; then
            echo ""
            echo "Démarrage du téléchargement de tous les modèles..."
            echo ""
            
            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"
            echo ""
            
            download_model \
                "https://huggingface.co/RunDiffusion/Juggernaut-XL-v9/resolve/main/Juggernaut-XL-v9.safetensors" \
                "$MODELS_DIR/Juggernaut-XL-v9.safetensors" \
                "Juggernaut XL v9" \
                "6.9GB"
            echo ""
            
            download_model \
                "https://huggingface.co/dathudehenry/ProteusV0.4/resolve/main/ProteusV0.4.safetensors" \
                "$MODELS_DIR/ProteusV0.4.safetensors" \
                "Proteus V0.4" \
                "6.9GB"
            echo ""
            
            download_model \
                "https://huggingface.co/Lykon/DreamShaper/resolve/main/DreamShaperXL1.0.safetensors" \
                "$MODELS_DIR/DreamShaperXL1.0.safetensors" \
                "DreamShaper XL 1.0" \
                "6.9GB"
            echo ""
            
            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"
            echo ""
            
            download_model \
                "https://huggingface.co/stabilityai/sdxl-vae/resolve/main/diffusion_pytorch_model.safetensors" \
                "$VAE_DIR/sdxl_vae.safetensors" \
                "SDXL VAE" \
                "335MB"
            
            clear
            print_header "✅ TOUS LES MODÈLES SONT TÉLÉCHARGÉS"
        else
            echo ""
            echo "Annulation..."
            exit 0
        fi
        ;;
    0)
        echo ""
        echo "Au revoir !"
        exit 0
        ;;
    *)
        echo ""
        echo -e "${RED}✗${NC} Choix invalide"
        exit 1
        ;;
esac

# Récapitulatif final
clear
print_header "✅ TÉLÉCHARGEMENT TERMINÉ"

echo "📊 Modèles installés:"
ls -lh "$MODELS_DIR"/*.safetensors 2>/dev/null | awk '{print "  • " $9 " (" $5 ")"}' || echo "  (aucun)"
echo ""
echo "📊 VAE installés:"
ls -lh "$VAE_DIR"/*.safetensors 2>/dev/null | awk '{print "  • " $9 " (" $5 ")"}' || echo "  (aucun)"
echo ""
echo -e "${GREEN}🚀 Prochaines étapes:${NC}"
echo ""
echo "1. Redémarrer ComfyUI:"
echo "   sudo docker restart comfyui"
echo ""
echo "2. Dans OpenWebUI (http://192.168.1.27:3000):"
echo "   Admin → Settings → Images"
echo ""
echo "3. Changer le nom du modèle:"
echo "   • RealVisXL_V4.0.safetensors"
echo "   • Juggernaut-XL-v9.safetensors"
echo "   • DreamShaperXL1.0.safetensors"
echo "   • sdxl_lightning_8step.safetensors"
echo ""
echo "4. Consulter la documentation:"
echo "   comfyui/docs/MODELS.md"
echo ""
echo "========================================"