#!/usr/bin/env python3
"""
Test script to verify the setup for HF Space
"""

def test_imports():
    """Test if all required modules can be imported"""
    print("🧪 Testing imports...")
    
    try:
        import gradio as gr
        print("✅ Gradio imported successfully")
    except ImportError as e:
        print(f"❌ Gradio import failed: {e}")
        return False
    
    try:
        import torch
        print(f"✅ PyTorch imported successfully (version: {torch.__version__})")
    except ImportError as e:
        print(f"❌ PyTorch import failed: {e}")
        return False
    
    try:
        import torchaudio
        print(f"✅ TorchAudio imported successfully")
    except ImportError as e:
        print(f"❌ TorchAudio import failed: {e}")
        return False
    
    return True

def test_hf_ac_imports():
    """Test hf_AC specific imports"""
    print("\n🔍 Testing hf_AC imports...")
    
    import sys
    from pathlib import Path
    
    # Add hf_AC to path
    current_dir = Path(__file__).parent
    hf_ac_path = current_dir / "hf_AC"
    
    if not hf_ac_path.exists():
        print("❌ hf_AC directory not found")
        return False
    
    sys.path.insert(0, str(hf_ac_path))
    
    try:
        from hf_AC.mmaudio.eval_utils import ModelConfig, all_model_cfg
        print("✅ hf_AC eval_utils imported successfully")
        
        if all_model_cfg:
            print(f"✅ Available model variants: {list(all_model_cfg.keys())}")
        else:
            print("⚠️ No model configurations found")
            
    except ImportError as e:
        print(f"❌ hf_AC import failed: {e}")
        return False
    
    return True

def test_gradio_interface():
    """Test basic Gradio interface creation"""
    print("\n🎨 Testing Gradio interface...")
    
    try:
        import gradio as gr
        
        def dummy_function(text):
            return f"Echo: {text}"
        
        # Create a simple interface
        iface = gr.Interface(
            fn=dummy_function,
            inputs=gr.Textbox(label="Test Input"),
            outputs=gr.Textbox(label="Test Output"),
            title="Test Interface"
        )
        
        print("✅ Gradio interface created successfully")
        return True
        
    except Exception as e:
        print(f"❌ Gradio interface creation failed: {e}")
        return False

def main():
    """Main test function"""
    print("🚀 Starting HF Space setup verification")
    print("=" * 50)
    
    # Test basic imports
    if not test_imports():
        print("\n❌ Basic import test failed")
        return False
    
    # Test hf_AC imports
    if not test_hf_ac_imports():
        print("\n⚠️ hf_AC import test failed (this is expected if model files are not downloaded)")
    
    # Test Gradio interface
    if not test_gradio_interface():
        print("\n❌ Gradio interface test failed")
        return False
    
    print("\n" + "=" * 50)
    print("✅ Setup verification completed successfully!")
    print("🎵 Your HF Space should be ready to deploy")
    
    return True

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