#!/usr/bin/env python3
"""
Test client for the OCR API
This script demonstrates how to use the OCR API endpoints
"""

import requests
import json
import sys
import os
from pathlib import Path
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# API base URL
BASE_URL = "http://localhost:8000"
API_KEY = os.getenv("API_KEY")

def test_health():
    """Test the health endpoint"""
    try:
        response = requests.get(f"{BASE_URL}/health")
        print(f"Health check: {response.status_code} - {response.json()}")
        return response.status_code == 200
    except requests.exceptions.ConnectionError:
        print("❌ Could not connect to the API. Make sure it's running on localhost:8000")
        return False

def test_pdf_ocr(file_path):
    """Test PDF OCR endpoint"""
    if not Path(file_path).exists():
        print(f"❌ File not found: {file_path}")
        return False

    try:
        with open(file_path, 'rb') as f:
            files = {'file': f}
            print(f"📄 Processing PDF: {file_path}")
            response = requests.post(f"{BASE_URL}/ocr/pdf", files=files)

        if response.status_code == 200:
            result = response.json()
            print(f"✅ PDF OCR successful!")
            print(f"   Pages: {result['pages']}")
            print(f"   Processing time: {result['processing_time']}s")
            print(f"   Confidence: {result['confidence']}%")
            print(f"   Text preview: {result['text'][:200]}...")
            return True
        else:
            print(f"❌ PDF OCR failed: {response.status_code} - {response.text}")
            return False

    except Exception as e:
        print(f"❌ Error processing PDF: {str(e)}")
        return False

def test_image_ocr(file_path):
    """Test image OCR endpoint"""
    if not Path(file_path).exists():
        print(f"❌ File not found: {file_path}")
        return False

    try:
        with open(file_path, 'rb') as f:
            files = {'file': f}
            print(f"🖼️  Processing image: {file_path}")
            response = requests.post(f"{BASE_URL}/ocr/image", files=files)

        if response.status_code == 200:
            result = response.json()
            print(f"✅ Image OCR successful!")
            print(f"   Processing time: {result['processing_time']}s")
            print(f"   Confidence: {result['confidence']}%")
            print(f"   Text: {result['text']}")
            return True
        else:
            print(f"❌ Image OCR failed: {response.status_code} - {response.text}")
            return False

    except Exception as e:
        print(f"❌ Error processing image: {str(e)}")
        return False

def main():
    """Main test function"""
    print("🧪 OCR API Test Client")
    print("=" * 50)

    # Test health endpoint
    if not test_health():
        sys.exit(1)

    print("\n" + "=" * 50)

    # Test PDF OCR if file provided
    if len(sys.argv) > 1:
        file_path = sys.argv[1]
        file_ext = Path(file_path).suffix.lower()

        if file_ext == '.pdf':
            test_pdf_ocr(file_path)
        elif file_ext in ['.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.tif']:
            test_image_ocr(file_path)
        else:
            print(f"❌ Unsupported file type: {file_ext}")
            print("Supported types: .pdf, .png, .jpg, .jpeg, .bmp, .tiff, .tif")
    else:
        print("📝 Usage: python test_client.py <file_path>")
        print("   Example: python test_client.py document.pdf")
        print("   Example: python test_client.py image.png")

if __name__ == "__main__":
    main()