#!/usr/bin/env python3
"""
Script pour pré-télécharger tous les modèles d'IA nécessaires.
Exécuter ce script avant la première utilisation pour éviter les délais.
"""

import sys
from pathlib import Path

# Ajouter le répertoire parent au path
sys.path.insert(0, str(Path(__file__).parent))

from core.config import CONFIG

def download_openclip():
    """Télécharge le modèle OpenCLIP pour les embeddings."""
    print("📥 Téléchargement d'OpenCLIP...")
    try:
        import open_clip
        model_name = CONFIG["models"]["clip"]
        model, _, preprocess = open_clip.create_model_and_transforms(
            model_name,
            pretrained="openai"
        )
        print(f"✅ OpenCLIP ({model_name}) téléchargé")
        return True
    except Exception as e:
        print(f"❌ Erreur OpenCLIP: {e}")
        return False

def download_owlv2():
    """Télécharge le modèle OWL-v2 pour la détection d'objets."""
    print("📥 Téléchargement d'OWL-v2...")
    try:
        from transformers import Owlv2Processor, Owlv2ForObjectDetection
        
        model_path = CONFIG["models"]["paths"]["owl_v2"]
        Path(model_path).mkdir(parents=True, exist_ok=True)
        
        processor = Owlv2Processor.from_pretrained(
            "google/owlv2-base-patch16-ensemble",
            cache_dir=model_path
        )
        model = Owlv2ForObjectDetection.from_pretrained(
            "google/owlv2-base-patch16-ensemble",
            cache_dir=model_path
        )
        print("✅ OWL-v2 téléchargé")
        return True
    except Exception as e:
        print(f"❌ Erreur OWL-v2: {e}")
        return False

def check_ollama():
    """Vérifie si Ollama est installé et le modèle Qwen est disponible."""
    print("🔍 Vérification d'Ollama...")
    try:
        import subprocess
        
        # Vérifier si ollama est installé
        result = subprocess.run(
            ["ollama", "list"],
            capture_output=True,
            text=True,
            timeout=5
        )
        
        if result.returncode == 0:
            print("✅ Ollama est installé")
            
            # Vérifier si qwen3-vl est disponible
            if "qwen3-vl" in result.stdout.lower():
                print("✅ Modèle Qwen3-VL disponible")
                return True
            else:
                print("⚠️  Modèle Qwen3-VL non trouvé")
                print("   Pour le télécharger: ollama pull qwen3-vl:8b")
                return False
        else:
            print("❌ Ollama n'est pas installé")
            print("   Installation: curl -fsSL https://ollama.com/install.sh | sh")
            return False
            
    except FileNotFoundError:
        print("❌ Ollama n'est pas installé")
        print("   Installation: curl -fsSL https://ollama.com/install.sh | sh")
        return False
    except Exception as e:
        print(f"❌ Erreur lors de la vérification d'Ollama: {e}")
        return False

def main():
    """Télécharge tous les modèles nécessaires."""
    print("=" * 60)
    print("🚀 Téléchargement des Modèles d'IA - RPA Vision V2")
    print("=" * 60)
    print()
    
    results = {
        "OpenCLIP": download_openclip(),
        "OWL-v2": download_owlv2(),
        "Ollama": check_ollama()
    }
    
    print()
    print("=" * 60)
    print("📊 Résumé")
    print("=" * 60)
    
    for model, success in results.items():
        status = "✅" if success else "❌"
        print(f"{status} {model}")
    
    print()
    
    if all(results.values()):
        print("🎉 Tous les modèles sont prêts !")
        print()
        print("Vous pouvez maintenant lancer l'application:")
        print("  ./run.sh")
        return 0
    else:
        print("⚠️  Certains modèles ne sont pas disponibles.")
        print()
        print("L'application fonctionnera en mode dégradé.")
        print("Consultez les messages ci-dessus pour plus d'informations.")
        return 1

if __name__ == "__main__":
    sys.exit(main())