#!/usr/bin/env python3
"""
Test script for Gradio app
"""

import sys
import os

def test_imports():
    """Test if all required modules can be imported"""
    print("Testing imports...")
    
    try:
        import gradio as gr
        print(f"✅ Gradio {gr.__version__}")
    except ImportError as e:
        print(f"❌ Gradio: {e}")
        return False
    
    try:
        import pandas as pd
        print(f"✅ Pandas {pd.__version__}")
    except ImportError as e:
        print(f"❌ Pandas: {e}")
        return False
    
    try:
        import numpy as np
        print(f"✅ NumPy {np.__version__}")
    except ImportError as e:
        print(f"❌ NumPy: {e}")
        return False
    
    return True

def test_config():
    """Test configuration import"""
    print("\nTesting configuration...")
    
    try:
        from config_hf import get_config, get_example_sequences, get_custom_css
        config = get_config()
        examples = get_example_sequences()
        css = get_custom_css()
        
        print("✅ Configuration loaded successfully")
        print(f"  - Max sequences: {config['MAX_SEQUENCES']}")
        print(f"  - Example sequences: {len(examples)} characters")
        print(f"  - CSS length: {len(css)} characters")
        return True
        
    except Exception as e:
        print(f"❌ Configuration failed: {e}")
        return False

def test_model_predictor():
    """Test model predictor"""
    print("\nTesting model predictor...")
    
    try:
        from model_predictor import EpitopePredictor
        predictor = EpitopePredictor()
        
        if predictor.model is None:
            print("⚠️ Model not loaded - will use demo mode")
        else:
            print("✅ Model loaded successfully")
        
        # Test prediction
        test_seq = "MKLLILTCLVAVALARPKHPIKHQGLPQEVLNENLLRFFVAPFPEVFGKEKVNEL"
        b_epitopes, t_epitopes = predictor.predict_epitopes(test_seq)
        
        print(f"✅ Prediction test: {len(b_epitopes)} B-cell, {len(t_epitopes)} T-cell epitopes")
        return True
        
    except Exception as e:
        print(f"❌ Model predictor failed: {e}")
        return False

def test_gradio_app():
    """Test Gradio app creation"""
    print("\nTesting Gradio app...")
    
    try:
        from app_gradio import create_interface
        demo = create_interface()
        
        print("✅ Gradio interface created successfully")
        print(f"  - Interface type: {type(demo)}")
        
        # Test if we can get the config
        if hasattr(demo, 'config'):
            print("✅ Interface has config")
        
        return True
        
    except Exception as e:
        print(f"❌ Gradio app creation failed: {e}")
        return False

def test_sample_prediction():
    """Test a sample prediction through the interface"""
    print("\nTesting sample prediction...")
    
    try:
        from app_gradio import predict_epitopes
        
        # Test with sample sequence
        sample_seq = """>Test_Protein
MKLLILTCLVAVALARPKHPIKHQGLPQEVLNENLLRFFVAPFPEVFGKEKVNEL"""
        
        # Mock progress function
        class MockProgress:
            def __call__(self, value, desc=""):
                print(f"  Progress: {value:.1%} - {desc}")
        
        result = predict_epitopes(sample_seq, None, 0.5, MockProgress())
        
        if len(result) == 5:  # Expected return format
            summary, b_df, t_df, csv_file, json_file = result
            print("✅ Prediction function works")
            print(f"  - Summary length: {len(summary)} characters")
            print(f"  - B-cell epitopes: {len(b_df)} rows")
            print(f"  - T-cell epitopes: {len(t_df)} rows")
            return True
        else:
            print(f"❌ Unexpected return format: {len(result)} items")
            return False
        
    except Exception as e:
        print(f"❌ Sample prediction failed: {e}")
        return False

def main():
    """Run all tests"""
    print("🧪 EpiPred Gradio App Test Suite")
    print("=" * 40)
    
    tests = [
        ("Basic Imports", test_imports),
        ("Configuration", test_config),
        ("Model Predictor", test_model_predictor),
        ("Gradio App", test_gradio_app),
        ("Sample Prediction", test_sample_prediction)
    ]
    
    passed = 0
    total = len(tests)
    
    for test_name, test_func in tests:
        try:
            if test_func():
                passed += 1
                print(f"✅ {test_name} PASSED")
            else:
                print(f"❌ {test_name} FAILED")
        except Exception as e:
            print(f"❌ {test_name} ERROR: {e}")
        print()
    
    print("=" * 40)
    print(f"Results: {passed}/{total} tests passed")
    
    if passed >= 4:  # Allow model predictor to fail
        print("🎉 Gradio app is ready for deployment!")
        if passed < total:
            print("⚠️ Some features may be limited (demo mode)")
    else:
        print("❌ App has issues that need to be fixed")
    
    return passed >= 4

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