import requests
import json
import time
from pathlib import Path

# Configuration
BASE_URL = "http://localhost:3015"
API_KEY = "jleu1212"
OCR_PDF_PATH = "ocr.pdf"

def test_ocr_upload_with_api_key():
    print("Testing OCR PDF upload with API key authentication...")
    print(f"API Key: {API_KEY}")
    print(f"OCR PDF: {OCR_PDF_PATH}")
    
    # Use query parameter for API key as per OpenAPI spec
    params = {"api_key_header_value": API_KEY}
    
    # Step 1: Check server status
    print("\n1. Checking server status...")
    try:
        status_response = requests.get(f"{BASE_URL}/", params=params)
        if status_response.status_code == 200:
            print("✓ Server is running and accessible")
        else:
            print(f"⚠ Server status: {status_response.status_code}")
    except Exception as e:
        print(f"✗ Server check error: {e}")
        return False
    
    # Step 2: Clear existing documents first
    print("\n2. Clearing existing documents...")
    try:
        clear_response = requests.delete(f"{BASE_URL}/documents", params=params)
        if clear_response.status_code == 200:
            clear_result = clear_response.json()
            print(f"✓ Documents cleared successfully: {clear_result}")
        else:
            print(f"⚠ Clear documents response: {clear_response.status_code} - {clear_response.text}")
    except Exception as e:
        print(f"⚠ Clear documents error: {e}")
    
    # Step 3: Upload OCR PDF
    print("\n3. Uploading OCR PDF...")
    
    if not Path(OCR_PDF_PATH).exists():
        print(f"✗ OCR PDF file not found: {OCR_PDF_PATH}")
        return False
    
    files = {"file": (OCR_PDF_PATH, open(OCR_PDF_PATH, "rb"), "application/pdf")}
    
    try:
        upload_response = requests.post(f"{BASE_URL}/documents/upload", files=files, params=params)
        print(f"Upload response status: {upload_response.status_code}")
        print(f"Upload response text: {upload_response.text}")
        
        if upload_response.status_code == 200:
            upload_result = upload_response.json()
            print(f"✓ Upload successful! Result: {json.dumps(upload_result, indent=2)}")
            
            # Step 4: Wait for processing and check document status
            print("\n4. Waiting for document processing...")
            time.sleep(10)  # Give more time for OCR processing
            
            try:
                status_response = requests.get(f"{BASE_URL}/documents", params=params)
                if status_response.status_code == 200:
                    documents = status_response.json()
                    print(f"✓ Documents status: {json.dumps(documents, indent=2)}")
                else:
                    print(f"⚠ Documents status check failed: {status_response.status_code}")
            except Exception as e:
                print(f"⚠ Documents status check error: {e}")
            
            # Step 5: Test search functionality
            print("\n5. Testing search functionality...")
            search_data = {
                "query": "safety precautions",
                "top_k": 5
            }
            
            try:
                search_response = requests.post(f"{BASE_URL}/query", json=search_data, params=params)
                if search_response.status_code == 200:
                    search_results = search_response.json()
                    print(f"✓ Search successful! Response: {json.dumps(search_results, indent=2)}")
                else:
                    print(f"⚠ Search failed: {search_response.status_code} - {search_response.text}")
            except Exception as e:
                print(f"⚠ Search error: {e}")
            
            return True
        else:
            print(f"✗ Upload failed: {upload_response.status_code} - {upload_response.text}")
            return False
            
    except Exception as e:
        print(f"✗ Upload error: {e}")
        return False
    finally:
        if 'files' in locals():
            files["file"][1].close()

if __name__ == "__main__":
    success = test_ocr_upload_with_api_key()
    if success:
        print("\n🎉 OCR upload test with API key authentication completed successfully!")
    else:
        print("\n❌ OCR upload test with API key authentication failed!")