import requests
import time
import os

def test_server_status():
    """Test if the server is running"""
    print("Testing server status...")
    try:
        response = requests.get("http://localhost:3015/health", timeout=5)
        if response.status_code == 200:
            print("✅ Server is running!")
            return True
        else:
            print(f"❌ Server returned status: {response.status_code}")
            return False
    except Exception as e:
        print(f"❌ Server not accessible: {e}")
        return False

def test_webui_access():
    """Test if web UI is accessible"""
    print("\nTesting web UI access...")
    try:
        response = requests.get("http://localhost:3015/webui/", timeout=5)
        if response.status_code == 200:
            print("✅ Web UI is accessible!")
            return True
        else:
            print(f"❌ Web UI returned status: {response.status_code}")
            return False
    except Exception as e:
        print(f"❌ Web UI not accessible: {e}")
        return False

def test_login():
    """Test login with credentials"""
    print("\nTesting login...")
    login_data = {
        "username": "jleu3482",
        "password": "jleu1212"
    }
    
    try:
        response = requests.post("http://localhost:3015/login", data=login_data, timeout=10)
        print(f"Login response: {response.status_code}")
        
        if response.status_code == 200:
            print("✅ Login successful!")
            data = response.json()
            print(f"   Token type: {data.get('token_type')}")
            print(f"   Auth mode: {data.get('auth_mode')}")
            return data.get('access_token')
        else:
            print(f"❌ Login failed: {response.text}")
            return None
    except Exception as e:
        print(f"❌ Login error: {e}")
        return None

def upload_ocr_pdf(token=None):
    """Upload the OCR PDF file for testing"""
    print("\nUploading OCR PDF...")
    
    headers = {}
    if token:
        headers['Authorization'] = f'Bearer {token}'
    
    files = {
        'file': ('ocr.pdf', open('../ocr.pdf', 'rb'), 'application/pdf')
    }
    
    try:
        response = requests.post(
            "http://localhost:3015/upload",
            files=files,
            headers=headers,
            timeout=30
        )
        
        print(f"Upload response: {response.status_code}")
        
        if response.status_code == 200:
            print("✅ Upload successful!")
            data = response.json()
            print(f"   Document ID: {data.get('doc_id')}")
            print(f"   Status: {data.get('status')}")
            return data.get('doc_id')
        else:
            print(f"❌ Upload failed: {response.text}")
            return None
    except Exception as e:
        print(f"❌ Upload error: {e}")
        return None
    finally:
        if 'files' in locals():
            files['file'][1].close()

def check_document_status(doc_id, token=None):
    """Check the status of the uploaded document"""
    print(f"\nChecking document status for ID: {doc_id}")
    
    headers = {}
    if token:
        headers['Authorization'] = f'Bearer {token}'
    
    try:
        response = requests.get(
            f"http://localhost:3015/doc-status/{doc_id}",
            headers=headers,
            timeout=10
        )
        
        print(f"Status response: {response.status_code}")
        
        if response.status_code == 200:
            data = response.json()
            print(f"✅ Document status: {data.get('status')}")
            print(f"   Progress: {data.get('progress', 0)}%")
            print(f"   Message: {data.get('message', '')}")
            return data
        else:
            print(f"❌ Status check failed: {response.text}")
            return None
    except Exception as e:
        print(f"❌ Status check error: {e}")
        return None

def test_search(query, token=None):
    """Test search functionality"""
    print(f"\nTesting search with query: '{query}'")
    
    headers = {'Content-Type': 'application/json'}
    if token:
        headers['Authorization'] = f'Bearer {token}'
    
    search_data = {
        "query": query,
        "top_k": 5
    }
    
    try:
        response = requests.post(
            "http://localhost:3015/query",
            json=search_data,
            headers=headers,
            timeout=15
        )
        
        print(f"Search response: {response.status_code}")
        
        if response.status_code == 200:
            data = response.json()
            print("✅ Search successful!")
            print(f"   Found {len(data.get('results', []))} results")
            
            for i, result in enumerate(data.get('results', [])):
                print(f"   Result {i+1}:")
                print(f"     Score: {result.get('score', 0):.3f}")
                print(f"     Text: {result.get('text', '')[:100]}...")
            
            return data
        else:
            print(f"❌ Search failed: {response.text}")
            return None
    except Exception as e:
        print(f"❌ Search error: {e}")
        return None

if __name__ == "__main__":
    print("LightRAG Server and OCR Test")
    print("=" * 50)
    
    # Test server status
    if not test_server_status():
        print("\n💡 Server is not running. Please start it with:")
        print("   cd LightRAG-main && .\\zrun.bat")
        exit(1)
    
    # Test web UI access
    test_webui_access()
    
    # Test login
    token = test_login()
    
    # Upload OCR PDF
    doc_id = upload_ocr_pdf(token)
    
    if doc_id:
        # Wait a bit for processing
        print("\n⏳ Waiting for document processing...")
        time.sleep(10)
        
        # Check document status
        status_data = check_document_status(doc_id, token)
        
        # Test search with OCR content
        test_search("optical character recognition", token)
        test_search("PDF document", token)
        test_search("text extraction", token)
    
    print("\n" + "=" * 50)
    print("Test completed!")
    print("Access the web UI at: http://localhost:3015/webui/")
    print("Username: jleu3482")
    print("Password: jleu1212")