"""
Quick test script for virus prediction API
Run this after starting the server to verify everything works
"""

import requests
import json

# API base URL (change if deployed)
BASE_URL = "http://localhost:7860"

def test_health():
    """Test health endpoint"""
    print("🔍 Testing /health endpoint...")
    response = requests.get(f"{BASE_URL}/health")
    print(f"Status: {response.status_code}")
    print(f"Response: {json.dumps(response.json(), indent=2)}\n")
    return response.status_code == 200

def test_info():
    """Test info endpoint"""
    print("🔍 Testing /info endpoint...")
    response = requests.get(f"{BASE_URL}/info")
    print(f"Status: {response.status_code}")
    print(f"Response: {json.dumps(response.json(), indent=2)}\n")
    return response.status_code == 200

def test_mappings():
    """Test mappings endpoint"""
    print("🔍 Testing /mappings endpoint...")
    response = requests.get(f"{BASE_URL}/mappings")
    print(f"Status: {response.status_code}")
    data = response.json()
    print(f"Total major viruses: {len(data.get('virus_mapping', {}))}")
    print(f"Total other viruses: {len(data.get('other_virus_mapping', {}))}")
    print(f"Total symptoms: {len(data.get('symptoms', []))}\n")
    return response.status_code == 200

def test_location_mappings():
    """Test location mappings endpoint for frontend configuration."""
    print("🔍 Testing /location-mappings endpoint...")
    response = requests.get(f"{BASE_URL}/location-mappings")
    print(f"Status: {response.status_code}")
    if response.status_code == 200:
        data = response.json()
        print(f"States: {len(data.get('states', []))}")
        print(f"Districts by state: {len(data.get('districts_by_state', {}))}")
        print(f"State mappings: {len(data.get('state_mapping', {}))}")
        print(f"District mappings: {len(data.get('district_mapping', {}))}")
        print(f"Source: {data.get('source', 'unknown')}\n")
    return response.status_code == 200

def test_prediction():
    """Test prediction endpoint with sample data"""
    print("🔍 Testing /predict endpoint...")
    
    # Sample patient data (Dengue-like symptoms)
    payload = {
        "age": 30.0,
        "SEX": 1,  # Male
        "PATIENTTYPE": 1,  # Inpatient
        "durationofillness": 3,
        "labstate": 32,  # Tamil Nadu
        "districtencoded": 120,
        "month": 8,  # August (monsoon)
        "year": 2024,
        "syndrome": 17,  # Fever >= 7 days
        # Dengue symptoms
        "FEVER": 1,
        "HEADACHE": 1,
        "MYALGIA": 1,
        "ARTHRALGIA": 1,
        "RETROORBITALPAIN": 1,
        "NAUSEA": 1,
        # All other symptoms 0
        "IRRITABILITY": 0,
        "ALTEREDSENSORIUM": 0,
        "SOMNOLENCE": 0,
        "NECKRIGIDITY": 0,
        "SEIZURES": 0,
        "DIARRHEA": 0,
        "DYSENTERY": 0,
        "VOMITING": 0,
        "ABDOMINALPAIN": 0,
        "MALAISE": 0,
        "CHILLS": 0,
        "RIGORS": 0,
        "BREATHLESSNESS": 0,
        "COUGH": 0,
        "RHINORRHEA": 0,
        "SORETHROAT": 0,
        "BULLAE": 0,
        "PAPULARRASH": 0,
        "PUSTULARRASH": 0,
        "MUSCULARRASH": 0,
        "MACULOPAPULARRASH": 0,
        "ESCHAR": 0,
        "DARKURINE": 0,
        "HEPATOMEGALY": 0,
        "JAUNDICE": 0,
        "REDEYE": 0,
        "DISCHARGEEYES": 0,
        "CRUSHINGEYES": 0,
        "SWELLINGEYES": 0
    }
    
    response = requests.post(f"{BASE_URL}/predict", json=payload)
    print(f"Status: {response.status_code}")
    
    if response.status_code == 200:
        data = response.json()
        print(f"\n✅ Prediction successful!")
        print(f"Predicted Virus: {data['predicted_virus']}")
        print(f"Confidence: {data['confidence']:.2f}%")
        print(f"\nTop 5 Predictions:")
        for i, pred in enumerate(data['top_5_predictions'], 1):
            print(f"  {i}. {pred['virus']}: {pred['confidence']:.2f}%")
        
        if data.get('sub_classification'):
            sub = data['sub_classification']
            print(f"\n🔍 Sub-classification:")
            print(f"  {sub['predicted_sub_virus']}: {sub['sub_confidence']:.2f}%")
        
        print(f"\nPrediction ID: {data.get('prediction_id', 'N/A')}")
        print(f"Timestamp: {data['timestamp']}\n")
        return True
    else:
        print(f"❌ Prediction failed")
        print(f"Error: {response.text}\n")
        return False

def test_stats():
    """Test stats endpoint"""
    print("🔍 Testing /stats endpoint...")
    response = requests.get(f"{BASE_URL}/stats")
    print(f"Status: {response.status_code}")
    if response.status_code == 200:
        print(f"Response: {json.dumps(response.json(), indent=2)}\n")
    return response.status_code == 200

def run_all_tests():
    """Run all tests"""
    print("=" * 60)
    print("🧪 VIRUS PREDICTION API - TEST SUITE")
    print("=" * 60)
    print()
    
    tests = [
        ("Health Check", test_health),
        ("API Info", test_info),
        ("Mappings", test_mappings),
        ("Location Mappings", test_location_mappings),
        ("Prediction", test_prediction),
        ("Statistics", test_stats),
    ]
    
    results = []
    for name, test_func in tests:
        try:
            success = test_func()
            results.append((name, "✅ PASS" if success else "❌ FAIL"))
        except requests.exceptions.ConnectionError:
            print(f"❌ Connection Error: Is the server running on {BASE_URL}?\n")
            results.append((name, "❌ CONNECTION ERROR"))
            break
        except Exception as e:
            print(f"❌ Error: {e}\n")
            results.append((name, f"❌ ERROR: {str(e)}"))
    
    print("=" * 60)
    print("📊 TEST RESULTS")
    print("=" * 60)
    for name, result in results:
        print(f"{name:.<40} {result}")
    print("=" * 60)

if __name__ == "__main__":
    run_all_tests()