#!/usr/bin/env python3
"""
Quick deployment test script for Simple AI Assistant
Run this to verify everything works before deploying to Hugging Face Spaces
"""

import sys
import importlib.util

def test_basic_imports():
    """Test if basic imports work"""
    print("🔍 Testing basic imports...")
    
    try:
        import torch
        print(f"✅ PyTorch {torch.__version__} imported successfully")
        
        # Check PyTorch version for security
        if torch.__version__ >= "2.6.0":
            print("✅ PyTorch version is secure (2.6.0+)")
        else:
            print(f"⚠️  PyTorch version {torch.__version__} may have security issues. Upgrade to 2.6.0+")
        
    except ImportError as e:
        print(f"❌ PyTorch import failed: {e}")
        return False
    
    try:
        import transformers
        print(f"✅ Transformers {transformers.__version__} imported successfully")
    except ImportError as e:
        print(f"❌ Transformers import failed: {e}")
        return False
    
    try:
        import gradio
        print(f"✅ Gradio {gradio.__version__} imported successfully")
    except ImportError as e:
        print(f"❌ Gradio import failed: {e}")
        return False
    
    try:
        import numpy
        print(f"✅ NumPy {numpy.__version__} imported successfully")
    except ImportError as e:
        print(f"❌ NumPy import failed: {e}")
        return False
    
    # Optional: Test autoawq (for Mistral model)
    try:
        import awq
        print(f"✅ AutoAWQ imported successfully")
    except ImportError:
        print("⚠️  AutoAWQ not available - Mistral model will fall back to DialoGPT")
    
    return True

def test_model_loading():
    """Test if we can load the Mistral-7B-AWQ model"""
    print("\n🤖 Testing Mistral-7B-AWQ model loading...")
    
    try:
        from transformers import AutoModelForCausalLM, AutoTokenizer
        
        # Test Mistral-7B-AWQ (as requested)
        model_id = "TheBloke/Mistral-7B-Instruct-v0.2-AWQ"
        print(f"🔄 Testing {model_id}...")
        
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        print("✅ Tokenizer loaded successfully")
        
        # Try AWQ model loading
        try:
            model = AutoModelForCausalLM.from_pretrained(
                model_id,
                device_map="auto" if torch.cuda.is_available() else "cpu",
                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                low_cpu_mem_usage=True,
                trust_remote_code=True,
                use_safetensors=True  # Secure loading
            )
            print("✅ AWQ model loaded successfully with device mapping!")
        except Exception as awq_error:
            print(f"⚠️ AWQ device mapping failed: {awq_error}")
            print("🔄 Trying standard loading...")
            model = AutoModelForCausalLM.from_pretrained(
                model_id,
                torch_dtype=torch.float32,  # Use float32 for compatibility
                low_cpu_mem_usage=True,
                trust_remote_code=True,
                use_safetensors=True
            )
            print("✅ AWQ model loaded successfully with standard loading!")
        
        # Test tokenization
        test_input = "Hello, how are you?"
        tokens = tokenizer.encode(test_input)
        print(f"✅ Tokenization test passed ({len(tokens)} tokens)")
        
        return True
        
    except Exception as e:
        print(f"❌ Mistral model loading failed: {e}")
        print("💡 Make sure autoawq>=0.1.8 is installed")
        return False

def test_emotion_detection():
    """Test emotion detection pipeline"""
    print("\n😊 Testing emotion detection...")
    
    try:
        from transformers import pipeline
        
        emotion_detector = pipeline(
            "sentiment-analysis",
            model="distilbert-base-uncased-finetuned-sst-2-english",
            return_all_scores=True
        )
        
        # Test emotion detection
        test_messages = [
            "I'm so happy today!",
            "I'm feeling really sad.",
            "The weather is okay."
        ]
        
        for msg in test_messages:
            result = emotion_detector(msg)
            print(f"✅ '{msg}' -> {result[0][0]['label']}")
        
        print("✅ Emotion detection working correctly")
        return True
        
    except Exception as e:
        print(f"❌ Emotion detection failed: {e}")
        return False

def test_gradio_interface():
    """Test if Gradio can create the interface"""
    print("\n🌐 Testing Gradio interface...")
    
    try:
        import gradio as gr
        
        # Test basic interface creation
        with gr.Blocks() as demo:
            gr.Markdown("# Test Interface")
            chatbot = gr.Chatbot()
            msg = gr.Textbox()
        
        print("✅ Gradio interface created successfully")
        print("✅ Ready for deployment!")
        return True
        
    except Exception as e:
        print(f"❌ Gradio interface test failed: {e}")
        return False

def main():
    """Run all tests"""
    print("🧪 Simple AI Assistant Deployment Test")
    print("=" * 50)
    
    all_passed = True
    
    # Run tests
    tests = [
        ("Basic Imports", test_basic_imports),
        ("Model Loading", test_model_loading),
        ("Emotion Detection", test_emotion_detection),
        ("Gradio Interface", test_gradio_interface)
    ]
    
    for test_name, test_func in tests:
        print(f"\n📋 Running {test_name} test...")
        try:
            if not test_func():
                all_passed = False
        except Exception as e:
            print(f"❌ {test_name} test crashed: {e}")
            all_passed = False
    
    print("\n" + "=" * 50)
    if all_passed:
        print("🎉 ALL TESTS PASSED! Your app is ready for deployment!")
        print("\n📋 Deployment Instructions:")
        print("1. Upload app.py and requirements.txt to Hugging Face Spaces")
        print("2. Set Space SDK to 'gradio'")
        print("3. Set Python version to 3.10+")
        print("4. Your app should build and run successfully!")
    else:
        print("❌ Some tests failed. Please fix the issues before deploying.")
        print("\n💡 Troubleshooting:")
        print("- Try using requirements_minimal.txt if main requirements fail")
        print("- Check Python version (needs 3.10+)")
        print("- Verify internet connection for model downloads")
        
    return all_passed

if __name__ == "__main__":
    # Allow importing torch in test
    try:
        import torch
    except ImportError:
        print("❌ PyTorch not installed. Please install requirements first:")
        print("pip install -r requirements.txt")
        sys.exit(1)
    
    success = main()
    sys.exit(0 if success else 1)