#!/usr/bin/env python3
"""
Test script for the RAG service
"""

import json
import requests
import time
import sys
from pathlib import Path

# Configuration
BASE_URL = "http://localhost:5002/rag-service"
TEST_DATA_FILE = Path(__file__).parent / "test_data.txt"
TEST_QUESTIONS_FILE = Path(__file__).parent / "test_questions.txt"

def test_health_check():
    """Test health check endpoint"""
    print("🔍 Testing health check...")
    try:
        response = requests.get(f"{BASE_URL}/health")
        if response.status_code == 200:
            data = response.json()
            print(f"✅ Health check passed: {data['status']}")
            return True
        else:
            print(f"❌ Health check failed: {response.status_code}")
            return False
    except Exception as e:
        print(f"❌ Health check error: {e}")
        return False

def test_ingest_content():
    """Test content ingestion"""
    print("📥 Testing content ingestion...")
    
    if not TEST_DATA_FILE.exists():
        print(f"❌ Test data file not found: {TEST_DATA_FILE}")
        return False
    
    try:
        with open(TEST_DATA_FILE, 'r') as f:
            content = f.read()
        
        payload = {
            "content": content,
            "clear_existing": True,
            "session_id": "test_session"
        }
        
        response = requests.post(f"{BASE_URL}/ingest", json=payload)
        
        if response.status_code == 200:
            data = response.json()
            print(f"✅ Content ingested successfully: {data['chunks_count']} chunks")
            return True
        else:
            print(f"❌ Content ingestion failed: {response.status_code} - {response.text}")
            return False
            
    except Exception as e:
        print(f"❌ Content ingestion error: {e}")
        return False

def test_query():
    """Test querying"""
    print("🔍 Testing query...")
    
    try:
        payload = {
            "query": "How do I register for a course?"
        }
        
        response = requests.post(f"{BASE_URL}/query", json=payload)
        
        if response.status_code == 200:
            data = response.json()
            print(f"✅ Query successful")
            print(f"📝 Answer: {data['answer'][:100]}...")
            print(f"📚 Sources: {len(data['sources'])} found")
            return True
        else:
            print(f"❌ Query failed: {response.status_code} - {response.text}")
            return False
            
    except Exception as e:
        print(f"❌ Query error: {e}")
        return False

def test_multiple_queries():
    """Test multiple queries from test file"""
    print("🔍 Testing multiple queries...")
    
    if not TEST_QUESTIONS_FILE.exists():
        print(f"❌ Test questions file not found: {TEST_QUESTIONS_FILE}")
        return False
    
    try:
        with open(TEST_QUESTIONS_FILE, 'r') as f:
            questions = f.read().split('\n')
        
        # Filter out empty lines and comments
        questions = [q.strip() for q in questions if q.strip() and not q.strip().startswith('#') and not q.strip().startswith('10 tricky')]
        
        success_count = 0
        for i, question in enumerate(questions[:5]):  # Test first 5 questions
            if question.startswith(tuple('0123456789')):
                # Remove numbering
                question = question.split('.', 1)[1].strip() if '.' in question else question
            
            print(f"  Query {i+1}: {question[:50]}...")
            
            payload = {"query": question}
            response = requests.post(f"{BASE_URL}/query", json=payload)
            
            if response.status_code == 200:
                success_count += 1
                print(f"    ✅ Success")
            else:
                print(f"    ❌ Failed: {response.status_code}")
            
            time.sleep(0.5)  # Brief delay between requests
        
        print(f"✅ Multiple queries test: {success_count}/{len(questions[:5])} successful")
        return success_count > 0
        
    except Exception as e:
        print(f"❌ Multiple queries error: {e}")
        return False

def test_status():
    """Test status endpoint"""
    print("📊 Testing status...")
    try:
        response = requests.get(f"{BASE_URL}/status")
        if response.status_code == 200:
            data = response.json()
            print(f"✅ Status check passed")
            print(f"   Documents: {data.get('documents', 0)}")
            print(f"   Chunks: {data.get('chunks', 0)}")
            print(f"   Model loaded: {data.get('embedding_model_loaded', False)}")
            print(f"   Gemini available: {data.get('gemini_available', False)}")
            return True
        else:
            print(f"❌ Status check failed: {response.status_code}")
            return False
    except Exception as e:
        print(f"❌ Status check error: {e}")
        return False

def main():
    """Run all tests"""
    print("🚀 Starting RAG Service Tests")
    print(f"🔗 Base URL: {BASE_URL}")
    print("-" * 50)
    
    tests = [
        ("Health Check", test_health_check),
        ("Content Ingestion", test_ingest_content),
        ("Status Check", test_status),
        ("Single Query", test_query),
        ("Multiple Queries", test_multiple_queries),
    ]
    
    results = []
    for test_name, test_func in tests:
        print(f"\n{'='*20} {test_name} {'='*20}")
        success = test_func()
        results.append((test_name, success))
        print(f"Result: {'✅ PASS' if success else '❌ FAIL'}")
    
    print("\n" + "="*60)
    print("📊 TEST SUMMARY")
    print("="*60)
    
    passed = sum(1 for _, success in results if success)
    total = len(results)
    
    for test_name, success in results:
        status = "✅ PASS" if success else "❌ FAIL"
        print(f"{test_name:20} | {status}")
    
    print("-" * 60)
    print(f"Total: {passed}/{total} tests passed")
    
    if passed == total:
        print("🎉 All tests passed!")
        sys.exit(0)
    else:
        print("❌ Some tests failed!")
        sys.exit(1)

if __name__ == "__main__":
    main()