#!/usr/bin/env python3
"""
Model Verification Script

This script tests all models used in the application to ensure they're working correctly.
"""

import sys
from typing import Dict, Any

def test_sentence_transformers():
    """Test sentence transformer models."""
    results = {}
    
    models_to_test = [
        "sentence-transformers/all-MiniLM-L6-v2",
        "BAAI/bge-m3"
    ]
    
    try:
        from sentence_transformers import SentenceTransformer
        
        for model_name in models_to_test:
            try:
                print(f"Testing {model_name}...")
                model = SentenceTransformer(model_name)
                embeddings = model.encode(["This is a test sentence."], show_progress_bar=False)
                
                if embeddings is not None and len(embeddings) > 0:
                    results[model_name] = "✅ PASS"
                    print(f"  ✅ {model_name} working correctly")
                else:
                    results[model_name] = "❌ FAIL - No embeddings generated"
                    print(f"  ❌ {model_name} failed to generate embeddings")
                    
            except Exception as e:
                results[model_name] = f"❌ FAIL - {str(e)}"
                print(f"  ❌ {model_name} failed: {e}")
                
    except ImportError:
        results["sentence-transformers"] = "❌ FAIL - Package not installed"
        print("❌ sentence-transformers package not installed")
    
    return results

def test_cross_encoder():
    """Test cross-encoder model."""
    results = {}
    model_name = "cross-encoder/ms-marco-MiniLM-L-6-v2"
    
    try:
        from sentence_transformers import CrossEncoder
        
        print(f"Testing {model_name}...")
        model = CrossEncoder(model_name)
        scores = model.predict([("test query", "test document")])
        
        if scores is not None and len(scores) > 0:
            results[model_name] = "✅ PASS"
            print(f"  ✅ {model_name} working correctly")
        else:
            results[model_name] = "❌ FAIL - No scores generated"
            print(f"  ❌ {model_name} failed to generate scores")
            
    except ImportError:
        results["cross-encoder"] = "❌ FAIL - sentence-transformers not installed"
        print("❌ sentence-transformers package not installed")
    except Exception as e:
        results[model_name] = f"❌ FAIL - {str(e)}"
        print(f"  ❌ {model_name} failed: {e}")
    
    return results

def test_transformers_pipeline():
    """Test transformers pipeline."""
    results = {}
    model_name = "MoritzLaurer/deberta-v3-base-zeroshot-v2.0"
    
    try:
        from transformers import pipeline
        
        print(f"Testing {model_name}...")
        classifier = pipeline(
            "zero-shot-classification",
            model=model_name,
            device=-1  # CPU
        )
        
        result = classifier(
            "This is a test sentence about policy.",
            ["policy", "technology", "sports"]
        )
        
        if result and 'labels' in result and len(result['labels']) > 0:
            results[model_name] = "✅ PASS"
            print(f"  ✅ {model_name} working correctly")
        else:
            results[model_name] = "❌ FAIL - No classification result"
            print(f"  ❌ {model_name} failed to classify")
            
    except ImportError:
        results["transformers"] = "❌ FAIL - transformers package not installed"
        print("❌ transformers package not installed")
    except Exception as e:
        results[model_name] = f"❌ FAIL - {str(e)}"
        print(f"  ❌ {model_name} failed: {e}")
    
    return results

def test_application_modules():
    """Test that application modules can be imported."""
    results = {}
    
    modules_to_test = [
        "utils.encoding_input",
        "utils.loading_embeddings", 
        "utils.retrieve_n_rerank",
        "utils.sentiment_analysis",
        "utils.coherence_bbscore",
        "utils.model_generation",
        "utils.generation_streaming"
    ]
    
    for module_name in modules_to_test:
        try:
            __import__(module_name)
            results[module_name] = "✅ PASS"
            print(f"✅ {module_name} imported successfully")
        except ImportError as e:
            results[module_name] = f"❌ FAIL - {str(e)}"
            print(f"❌ {module_name} import failed: {e}")
        except Exception as e:
            results[module_name] = f"❌ FAIL - {str(e)}"
            print(f"❌ {module_name} error: {e}")
    
    return results

def main():
    """Run all tests."""
    print("🧪 Model Verification Test Suite")
    print("=" * 50)
    
    all_results = {}
    
    print("\n📦 Testing Sentence Transformers...")
    all_results.update(test_sentence_transformers())
    
    print("\n🔄 Testing Cross Encoder...")
    all_results.update(test_cross_encoder())
    
    print("\n🤖 Testing Transformers Pipeline...")
    all_results.update(test_transformers_pipeline())
    
    print("\n📚 Testing Application Modules...")
    all_results.update(test_application_modules())
    
    # Summary
    print("\n" + "=" * 50)
    print("📋 TEST SUMMARY")
    print("=" * 50)
    
    passed = 0
    failed = 0
    
    for name, result in all_results.items():
        print(f"{result} {name}")
        if "✅ PASS" in result:
            passed += 1
        else:
            failed += 1
    
    print(f"\n📊 Results: {passed} passed, {failed} failed")
    
    if failed == 0:
        print("🎉 All tests passed! The application is ready to deploy.")
        return 0
    else:
        print("⚠️  Some tests failed. Please check the errors above.")
        return 1

if __name__ == "__main__":
    sys.exit(main())