#!/usr/bin/env python3
"""
Test script for the Text-to-SQL application
"""

import requests
import json
import time

def test_health():
    """Test health endpoint"""
    try:
        response = requests.get("http://localhost:8000/health")
        print(f"Health check: {response.status_code}")
        if response.status_code == 200:
            data = response.json()
            print(f"Status: {data['status']}")
            print(f"Model loaded: {data['model_loaded']}")
        return response.status_code == 200
    except Exception as e:
        print(f"Health check failed: {e}")
        return False

def test_single_prediction():
    """Test single prediction endpoint"""
    try:
        data = {
            "question": "How many employees are older than 30?",
            "table_headers": ["id", "name", "age", "department", "salary"]
        }
        
        response = requests.post("http://localhost:8000/predict", json=data)
        print(f"Single prediction: {response.status_code}")
        
        if response.status_code == 200:
            result = response.json()
            print(f"Question: {result['question']}")
            print(f"SQL: {result['sql_query']}")
            print(f"Processing time: {result['processing_time']:.3f}s")
            return True
        else:
            print(f"Error: {response.text}")
            return False
            
    except Exception as e:
        print(f"Single prediction failed: {e}")
        return False

def test_batch_prediction():
    """Test batch prediction endpoint"""
    try:
        data = {
            "queries": [
                {
                    "question": "How many employees are older than 30?",
                    "table_headers": ["id", "name", "age", "department", "salary"]
                },
                {
                    "question": "Show all employees in IT department",
                    "table_headers": ["id", "name", "age", "department", "salary"]
                }
            ]
        }
        
        response = requests.post("http://localhost:8000/batch", json=data)
        print(f"Batch prediction: {response.status_code}")
        
        if response.status_code == 200:
            result = response.json()
            print(f"Total queries: {result['total_queries']}")
            print(f"Successful queries: {result['successful_queries']}")
            
            for i, res in enumerate(result['results']):
                print(f"\nQuery {i+1}:")
                print(f"  Question: {res['question']}")
                print(f"  SQL: {res['sql_query']}")
            return True
        else:
            print(f"Error: {response.text}")
            return False
            
    except Exception as e:
        print(f"Batch prediction failed: {e}")
        return False

def main():
    """Run all tests"""
    print("🧪 Testing Text-to-SQL Application")
    print("=" * 50)
    
    # Wait a bit for the server to start
    print("Waiting for server to be ready...")
    time.sleep(5)
    
    # Test health
    print("\n1. Testing health endpoint...")
    health_ok = test_health()
    
    if not health_ok:
        print("❌ Health check failed. Make sure the server is running.")
        return
    
    # Test single prediction
    print("\n2. Testing single prediction...")
    single_ok = test_single_prediction()
    
    # Test batch prediction
    print("\n3. Testing batch prediction...")
    batch_ok = test_batch_prediction()
    
    # Summary
    print("\n" + "=" * 50)
    print("📊 Test Results:")
    print(f"Health check: {'✅' if health_ok else '❌'}")
    print(f"Single prediction: {'✅' if single_ok else '❌'}")
    print(f"Batch prediction: {'✅' if batch_ok else '❌'}")
    
    if all([health_ok, single_ok, batch_ok]):
        print("\n🎉 All tests passed! Your application is ready for deployment.")
    else:
        print("\n⚠️  Some tests failed. Please check the errors above.")

if __name__ == "__main__":
    main()