#!/usr/bin/env python3
"""
Script para verificar la configuración del Space NTIA
"""

import os
import sys
from huggingface_hub import login, model_info

def check_hf_token():
    """Verificar si el token de Hugging Face está configurado"""
    token = os.getenv("HF_TOKEN")
    if not token:
        print("❌ HF_TOKEN no está configurado")
        print("💡 Configura la variable de entorno HF_TOKEN en el Space")
        return False
    
    print(f"✅ HF_TOKEN encontrado: {token[:10]}...{token[-10:] if len(token) > 20 else '***'}")
    return True

def check_flux_access():
    """Verificar acceso a modelos FLUX"""
    token = os.getenv("HF_TOKEN")
    if not token:
        print("❌ No se puede verificar acceso sin HF_TOKEN")
        return False
    
    try:
        # Intentar acceder a FLUX.1-dev
        print("🔍 Verificando acceso a FLUX.1-dev...")
        info = model_info("black-forest-labs/FLUX.1-dev", token=token)
        print(f"✅ Acceso a FLUX.1-dev: {info.modelId}")
        
        # Intentar acceder a FLUX.1-schnell
        print("🔍 Verificando acceso a FLUX.1-schnell...")
        info = model_info("black-forest-labs/FLUX.1-schnell", token=token)
        print(f"✅ Acceso a FLUX.1-schnell: {info.modelId}")
        
        return True
        
    except Exception as e:
        print(f"❌ Error verificando acceso a FLUX: {e}")
        print("💡 Asegúrate de:")
        print("   1. Tener acceso a los modelos FLUX en Hugging Face")
        print("   2. Que el token tenga permisos de lectura")
        print("   3. Haber aceptado los términos de licencia")
        return False

def check_dependencies():
    """Verificar dependencias necesarias"""
    try:
        import torch
        print(f"✅ PyTorch: {torch.__version__}")
        
        import diffusers
        print(f"✅ Diffusers: {diffusers.__version__}")
        
        from diffusers import FluxPipeline
        print("✅ FluxPipeline disponible")
        
        return True
        
    except ImportError as e:
        print(f"❌ Dependencia faltante: {e}")
        return False

def main():
    """Función principal"""
    print("🔧 Verificando configuración del Space NTIA...")
    print("=" * 50)
    
    # Verificar token
    token_ok = check_hf_token()
    print()
    
    # Verificar dependencias
    deps_ok = check_dependencies()
    print()
    
    # Verificar acceso a FLUX
    if token_ok:
        flux_ok = check_flux_access()
    else:
        flux_ok = False
    print()
    
    # Resumen
    print("=" * 50)
    print("📊 RESUMEN DE CONFIGURACIÓN:")
    print(f"   Token HF: {'✅' if token_ok else '❌'}")
    print(f"   Dependencias: {'✅' if deps_ok else '❌'}")
    print(f"   Acceso FLUX: {'✅' if flux_ok else '❌'}")
    
    if token_ok and deps_ok and flux_ok:
        print("\n🎉 ¡Todo configurado correctamente!")
        print("   Los modelos FLUX deberían funcionar en el Space")
    else:
        print("\n⚠️ Hay problemas de configuración:")
        if not token_ok:
            print("   - Configura HF_TOKEN en las variables de entorno del Space")
        if not deps_ok:
            print("   - Instala las dependencias necesarias")
        if not flux_ok:
            print("   - Solicita acceso a los modelos FLUX en Hugging Face")

if __name__ == "__main__":
    main()