#!/usr/bin/env python3
"""
Script de configuración para Hugging Face Space
Resuelve problemas comunes de configuración y dependencias
"""

import os
import sys
import subprocess
import logging

# Configurar logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def check_environment():
    """Verifica el entorno de ejecución"""
    logger.info("🔍 Verificando entorno...")
    
    # Verificar variables de entorno del Space
    space_id = os.getenv("SPACE_ID")
    space_host = os.getenv("SPACE_HOST")
    
    if space_id:
        logger.info(f"✅ Space ID: {space_id}")
    if space_host:
        logger.info(f"✅ Space Host: {space_host}")
    
    # Verificar GPU
    try:
        import torch
        if torch.cuda.is_available():
            logger.info(f"✅ GPU disponible: {torch.cuda.get_device_name()}")
            logger.info(f"✅ Memoria GPU: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
        else:
            logger.warning("⚠️ GPU no disponible")
    except ImportError:
        logger.error("❌ PyTorch no está instalado")

def install_mmaudio():
    """Instala el paquete mmaudio local"""
    logger.info("📦 Instalando mmaudio...")
    
    try:
        # Verificar si mmaudio ya está instalado
        import mmaudio
        logger.info("✅ mmaudio ya está instalado")
        return True
    except ImportError:
        pass
    
    try:
        # Instalar mmaudio desde el directorio local
        result = subprocess.run(
            ["pip", "install", "-e", "."],
            capture_output=True,
            text=True,
            cwd=os.getcwd()
        )
        
        if result.returncode == 0:
            logger.info("✅ mmaudio instalado correctamente")
            return True
        else:
            logger.error(f"❌ Error instalando mmaudio: {result.stderr}")
            return False
    except Exception as e:
        logger.error(f"❌ Error en instalación: {e}")
        return False

def verify_dependencies():
    """Verifica que todas las dependencias estén instaladas"""
    logger.info("🔍 Verificando dependencias...")
    
    dependencies = [
        "torch",
        "torchvision", 
        "torchaudio",
        "diffusers",
        "transformers",
        "accelerate",
        "gradio",
        "numpy",
        "Pillow",
        "librosa",
        "soundfile",
        "av"
    ]
    
    missing = []
    for dep in dependencies:
        try:
            __import__(dep)
            logger.info(f"✅ {dep}")
        except ImportError:
            logger.error(f"❌ {dep} no está instalado")
            missing.append(dep)
    
    if missing:
        logger.warning(f"⚠️ Dependencias faltantes: {', '.join(missing)}")
        return False
    
    logger.info("✅ Todas las dependencias están instaladas")
    return True

def setup_models():
    """Configura los modelos necesarios"""
    logger.info("🤖 Configurando modelos...")
    
    # Verificar que los directorios de modelos existan
    model_dirs = ["models", "checkpoints"]
    for dir_name in model_dirs:
        if not os.path.exists(dir_name):
            os.makedirs(dir_name)
            logger.info(f"📁 Creado directorio: {dir_name}")

def main():
    """Función principal de configuración"""
    logger.info("🚀 Iniciando configuración del Space...")
    
    # Verificar entorno
    check_environment()
    
    # Instalar mmaudio
    if not install_mmaudio():
        logger.error("❌ No se pudo instalar mmaudio")
        return False
    
    # Verificar dependencias
    if not verify_dependencies():
        logger.error("❌ Faltan dependencias")
        return False
    
    # Configurar modelos
    setup_models()
    
    logger.info("✅ Configuración completada exitosamente")
    return True

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)