#!/usr/bin/env python3
"""
Test LightRAG authentication methods and upload OCR PDF
"""

import requests
import json
import base64

# Configuration
BASE_URL = "http://localhost:3015"
USERNAME = "jleu3482"
PASSWORD = "jleu1212"

def test_auth_methods():
    """Test different authentication methods"""
    
    print("=== TESTING AUTHENTICATION METHODS ===")
    
    # Method 1: Basic Auth with requests.auth
    print("\n1. Testing Basic Auth with requests.auth...")
    try:
        response = requests.get(f"{BASE_URL}/documents", auth=(USERNAME, PASSWORD))
        print(f"   Status: {response.status_code}")
        if response.status_code == 200:
            print("   ✅ Basic Auth successful!")
            return "basic_auth"
        else:
            print(f"   ❌ Failed: {response.text}")
    except Exception as e:
        print(f"   ❌ Error: {e}")
    
    # Method 2: Basic Auth with headers
    print("\n2. Testing Basic Auth with headers...")
    try:
        credentials = base64.b64encode(f"{USERNAME}:{PASSWORD}".encode()).decode()
        headers = {"Authorization": f"Basic {credentials}"}
        response = requests.get(f"{BASE_URL}/documents", headers=headers)
        print(f"   Status: {response.status_code}")
        if response.status_code == 200:
            print("   ✅ Basic Auth with headers successful!")
            return "basic_auth_headers"
        else:
            print(f"   ❌ Failed: {response.text}")
    except Exception as e:
        print(f"   ❌ Error: {e}")
    
    # Method 3: API Key
    print("\n3. Testing API Key...")
    try:
        headers = {"X-API-Key": "jleu1212"}
        response = requests.get(f"{BASE_URL}/documents", headers=headers)
        print(f"   Status: {response.status_code}")
        if response.status_code == 200:
            print("   ✅ API Key successful!")
            return "api_key"
        else:
            print(f"   ❌ Failed: {response.text}")
    except Exception as e:
        print(f"   ❌ Error: {e}")
    
    # Method 4: Bearer token (if available)
    print("\n4. Testing Bearer Token...")
    try:
        headers = {"Authorization": "Bearer jleu1212"}
        response = requests.get(f"{BASE_URL}/documents", headers=headers)
        print(f"   Status: {response.status_code}")
        if response.status_code == 200:
            print("   ✅ Bearer Token successful!")
            return "bearer"
        else:
            print(f"   ❌ Failed: {response.text}")
    except Exception as e:
        print(f"   ❌ Error: {e}")
    
    return None

def upload_with_auth(method):
    """Upload OCR PDF using the successful authentication method"""
    print(f"\n=== UPLOADING OCR.PDF WITH {method.upper()} ===")
    
    pdf_path = "ocr.pdf"
    headers = {}
    
    if method == "basic_auth":
        auth = (USERNAME, PASSWORD)
    elif method == "basic_auth_headers":
        credentials = base64.b64encode(f"{USERNAME}:{PASSWORD}".encode()).decode()
        headers = {"Authorization": f"Basic {credentials}"}
        auth = None
    elif method == "api_key":
        headers = {"X-API-Key": "jleu1212"}
        auth = None
    elif method == "bearer":
        headers = {"Authorization": "Bearer jleu1212"}
        auth = None
    else:
        print("❌ Unknown authentication method")
        return False
    
    try:
        with open(pdf_path, 'rb') as file:
            files = {'file': ('ocr.pdf', file, 'application/pdf')}
            
            if auth:
                response = requests.post(f"{BASE_URL}/documents/upload", files=files, auth=auth)
            else:
                response = requests.post(f"{BASE_URL}/documents/upload", files=files, headers=headers)
        
        print(f"Upload response: {response.status_code}")
        
        if response.status_code == 200:
            print("✅ OCR PDF uploaded successfully!")
            result = response.json()
            print(f"Document ID: {result.get('id', 'Unknown')}")
            return True
        else:
            print(f"❌ Upload failed: {response.text}")
            return False
            
    except Exception as e:
        print(f"❌ Upload error: {e}")
        return False

def main():
    print("Testing LightRAG Authentication and OCR Upload")
    print("=" * 50)
    
    # Test authentication methods
    auth_method = test_auth_methods()
    
    if auth_method:
        print(f"\n✅ Found working authentication method: {auth_method}")
        
        # Upload OCR PDF
        if upload_with_auth(auth_method):
            print("\n🎉 Success! OCR PDF uploaded successfully.")
            print("\nNext steps:")
            print("1. Wait for document processing to complete")
            print("2. Test search functionality")
            print("3. Access Web UI at: http://localhost:3015/webui/")
        else:
            print("\n❌ Failed to upload OCR PDF")
    else:
        print("\n❌ No authentication method worked")
        print("\nTroubleshooting steps:")
        print("1. Check if LightRAG server is running")
        print("2. Verify authentication configuration in .env file")
        print("3. Check server logs for authentication errors")

if __name__ == "__main__":
    main()