#!/usr/bin/env python3
"""
Test script for the Question Generation API
Run this after your Space is deployed to test the API endpoints
"""

import requests
import json
import time

# Replace with your actual Space URL
BASE_URL = "https://your-space-name.hf.space"

def test_health_endpoint():
    """Test the health check endpoint"""
    print("🔍 Testing health endpoint...")
    
    try:
        response = requests.get(f"{BASE_URL}/health", timeout=30)
        print(f"Status Code: {response.status_code}")
        
        if response.status_code == 200:
            data = response.json()
            print(f"✅ Health Check Passed")
            print(f"Model Loaded: {data['model_loaded']}")
            print(f"Device: {data['device']}")
            if data.get('memory_usage'):
                memory = data['memory_usage']
                print(f"VRAM Usage: {memory.get('allocated_gb', 0):.2f}GB / {memory.get('total_gb', 0):.2f}GB")
            return True
        else:
            print(f"❌ Health Check Failed: {response.text}")
            return False
            
    except requests.exceptions.RequestException as e:
        print(f"❌ Health Check Error: {e}")
        return False

def test_question_generation():
    """Test the question generation endpoint"""
    print("\n🤔 Testing question generation...")
    
    test_cases = [
        {
            "name": "Simple Statement",
            "data": {
                "statement": "Artificial intelligence is transforming healthcare by enabling more accurate diagnoses, personalized treatments, and efficient drug discovery processes.",
                "num_questions": 3,
                "difficulty_level": "medium"
            }
        },
        {
            "name": "Complex Statement",
            "data": {
                "statement": "Climate change represents one of the most significant challenges of the 21st century, involving complex interactions between atmospheric chemistry, ocean currents, biodiversity loss, and human economic systems. The greenhouse effect, primarily driven by carbon dioxide emissions from fossil fuel combustion, is causing global temperatures to rise at an unprecedented rate.",
                "num_questions": 5,
                "difficulty_level": "hard",
                "temperature": 0.9
            }
        },
        {
            "name": "Short Statement",
            "data": {
                "statement": "Water boils at 100 degrees Celsius at sea level.",
                "num_questions": 2,
                "difficulty_level": "easy"
            }
        }
    ]
    
    for i, test_case in enumerate(test_cases, 1):
        print(f"\n📝 Test Case {i}: {test_case['name']}")
        print(f"Statement: {test_case['data']['statement'][:100]}...")
        
        try:
            response = requests.post(
                f"{BASE_URL}/generate-questions",
                json=test_case['data'],
                timeout=60  # Increased timeout for model inference
            )
            
            print(f"Status Code: {response.status_code}")
            
            if response.status_code == 200:
                data = response.json()
                questions = data['questions']
                
                print(f"✅ Generated {len(questions)} questions:")
                for j, question in enumerate(questions, 1):
                    print(f"   {j}. {question}")
                
                print(f"Metadata: {data['metadata']}")
                
            else:
                print(f"❌ Generation Failed: {response.text}")
                
        except requests.exceptions.RequestException as e:
            print(f"❌ Request Error: {e}")

def test_error_handling():
    """Test error handling"""
    print("\n🚨 Testing error handling...")
    
    # Test invalid parameters
    invalid_tests = [
        {
            "name": "Missing statement",
            "data": {"num_questions": 3}
        },
        {
            "name": "Invalid num_questions",
            "data": {
                "statement": "Test statement",
                "num_questions": 15  # Too high
            }
        },
        {
            "name": "Invalid temperature",
            "data": {
                "statement": "Test statement",
                "temperature": 5.0  # Too high
            }
        }
    ]
    
    for test in invalid_tests:
        print(f"\n🔍 Testing: {test['name']}")
        try:
            response = requests.post(
                f"{BASE_URL}/generate-questions",
                json=test['data'],
                timeout=30
            )
            
            if response.status_code == 422:
                print("✅ Correctly rejected invalid input")
            else:
                print(f"⚠️ Unexpected status code: {response.status_code}")
                
        except requests.exceptions.RequestException as e:
            print(f"❌ Request Error: {e}")

def benchmark_performance():
    """Simple performance benchmark"""
    print("\n⚡ Performance Benchmark...")
    
    statement = "Machine learning algorithms are becoming increasingly sophisticated, enabling computers to learn patterns from data without being explicitly programmed for every scenario."
    
    times = []
    for i in range(3):
        print(f"Run {i+1}/3...", end=" ")
        
        start_time = time.time()
        try:
            response = requests.post(
                f"{BASE_URL}/generate-questions",
                json={
                    "statement": statement,
                    "num_questions": 3,
                    "difficulty_level": "medium"
                },
                timeout=60
            )
            
            end_time = time.time()
            duration = end_time - start_time
            times.append(duration)
            
            if response.status_code == 200:
                print(f"✅ {duration:.2f}s")
            else:
                print(f"❌ Failed ({response.status_code})")
                
        except requests.exceptions.RequestException as e:
            print(f"❌ Error: {e}")
    
    if times:
        avg_time = sum(times) / len(times)
        print(f"\n📊 Average Response Time: {avg_time:.2f}s")
        print(f"📊 Min: {min(times):.2f}s, Max: {max(times):.2f}s")

def main():
    """Run all tests"""
    print("🚀 Starting API Tests")
    print(f"Base URL: {BASE_URL}")
    print("=" * 50)
    
    # Test health first
    if not test_health_endpoint():
        print("\n❌ Health check failed. Make sure your Space is running and accessible.")
        return
    
    # Wait a moment for model to be ready
    print("\n⏳ Waiting for model to be ready...")
    time.sleep(5)
    
    # Run tests
    test_question_generation()
    test_error_handling()
    benchmark_performance()
    
    print("\n" + "=" * 50)
    print("✅ All tests completed!")
    print("\n💡 Usage Examples:")
    print(f"curl -X POST '{BASE_URL}/generate-questions' \\")
    print("     -H 'Content-Type: application/json' \\")
    print("     -d '{\"statement\": \"Your statement here\", \"num_questions\": 3}'")

if __name__ == "__main__":
    # Update this with your actual Space URL before running
    if "your-space-name" in BASE_URL:
        print("⚠️  Please update BASE_URL with your actual Space URL before running tests!")
        print("Example: BASE_URL = 'https://username-question-generation-api.hf.space'")
    else:
        main()