import requests
import json
import time

def test_ocr_search():
    # Test login and search functionality
    base_url = 'http://localhost:3015'
    login_data = {'username': 'jleu3482', 'password': 'jleu1212'}
    
    print("Testing LightRAG OCR search functionality...")
    
    try:
        # Login
        response = requests.post(f"{base_url}/login", data=login_data)
        print(f'Login response: {response.status_code}')
        
        if response.status_code == 200:
            print('Login successful')
            
            # Test search queries on OCR content
            search_queries = [
                "minimum safe distance",
                "safety precautions", 
                "high voltage",
                "traction voltage",
                "electrical installation"
            ]
            
            for query in search_queries:
                print(f'\nTesting search query: "{query}"')
                
                # Perform search
                search_data = {
                    'query': query,
                    'top_k': 5,
                    'rerank': False
                }
                
                search_response = requests.post(f"{base_url}/search", json=search_data)
                print(f'Search response: {search_response.status_code}')
                
                if search_response.status_code == 200:
                    results = search_response.json()
                    
                    if 'results' in results:
                        print(f'Found {len(results["results"])} results:')
                        for i, result in enumerate(results['results']):
                            print(f'  {i+1}. Score: {result.get("score", "N/A")}')
                            print(f'     Content: {result.get("content", "N/A")[:200]}...')
                            print(f'     Source: {result.get("source", "N/A")}')
                            print()
                    else:
                        print(f'No results found for query: {query}')
                else:
                    print(f'Search failed: {search_response.text}')
                
                time.sleep(1)  # Small delay between queries
            
            return True
        else:
            print(f'Login failed: {response.text}')
            return False
            
    except Exception as e:
        print(f'Error during search test: {e}')
        return False

if __name__ == '__main__':
    success = test_ocr_search()
    if success:
        print('\n✅ OCR search test PASSED!')
    else:
        print('\n❌ OCR search test FAILED!')