#!/usr/bin/env python3
"""
Test script for Hugging Face Dots-OCR Space API

This script demonstrates how to interact with your deployed Dots-OCR Space
at https://algoryn-dots-ocr-idcard.hf.space
"""

import requests
import json
import time
from pathlib import Path
from typing import Optional, Dict, Any, List


class HFDotsOCRClient:
    """Client for interacting with the Hugging Face Dots-OCR Space API."""
    
    def __init__(self, base_url: str = "https://algoryn-dots-ocr-idcard.hf.space"):
        """Initialize the client with the Space URL."""
        self.base_url = base_url.rstrip('/')
        self.session = requests.Session()
        # Set a reasonable timeout for HF Spaces
        self.session.timeout = 60
    
    def health_check(self) -> Dict[str, Any]:
        """Check if the Space is healthy and running."""
        try:
            response = self.session.get(f"{self.base_url}/health")
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            return {"error": f"Health check failed: {e}"}
    
    def extract_text(
        self, 
        image_path: str, 
        roi: Optional[Dict[str, float]] = None
    ) -> Dict[str, Any]:
        """Extract text from an identity document image.
        
        Args:
            image_path: Path to the image file
            roi: Optional region of interest as dict with x1, y1, x2, y2 (0-1 normalized)
        """
        try:
            with open(image_path, 'rb') as f:
                files = {'file': f}
                data = {}
                
                # Add ROI if provided
                if roi:
                    data['roi'] = json.dumps(roi)
                
                response = self.session.post(
                    f"{self.base_url}/v1/id/ocr",
                    files=files,
                    data=data
                )
                response.raise_for_status()
                return response.json()
                
        except requests.exceptions.RequestException as e:
            return {"error": f"OCR extraction failed: {e}"}
        except FileNotFoundError:
            return {"error": f"Image file not found: {image_path}"}


def print_ocr_results(result: Dict[str, Any]) -> None:
    """Pretty print OCR extraction results."""
    if "error" in result:
        print(f"❌ Error: {result['error']}")
        return
    
    print(f"✅ Request ID: {result.get('request_id', 'N/A')}")
    print(f"📊 Media Type: {result.get('media_type', 'N/A')}")
    print(f"⏱️  Processing Time: {result.get('processing_time', 0):.2f}s")
    
    detections = result.get('detections', [])
    print(f"🔍 OCR Detections: {len(detections)}")
    
    for i, detection in enumerate(detections, 1):
        print(f"\n📄 Detection {i}:")
        
        # Print MRZ data if available
        mrz_data = detection.get('mrz_data')
        if mrz_data:
            print(f"   🆔 MRZ Data:")
            print(f"     Format: {mrz_data.get('format_type', 'N/A')}")
            print(f"     Valid: {mrz_data.get('is_valid', False)}")
            print(f"     Confidence: {mrz_data.get('confidence', 0):.3f}")
            if mrz_data.get('raw_text'):
                print(f"     Raw Text: {mrz_data['raw_text'][:50]}...")
        else:
            print(f"   🆔 MRZ Data: None detected")
        
        # Print extracted fields
        extracted_fields = detection.get('extracted_fields', {})
        print(f"   📋 Extracted Fields:")
        
        # Define field categories for better organization
        field_categories = {
            "Document Info": [
                "document_number", "document_type", "issuing_country", "issuing_authority"
            ],
            "Personal Info": [
                "surname", "given_names", "nationality", "date_of_birth", 
                "gender", "place_of_birth"
            ],
            "Validity Info": [
                "date_of_issue", "date_of_expiry", "personal_number"
            ],
            "Additional": [
                "optional_data_1", "optional_data_2"
            ]
        }
        
        for category, fields in field_categories.items():
            category_fields = []
            for field_name in fields:
                field_data = extracted_fields.get(field_name)
                if field_data and field_data.get('value'):
                    category_fields.append(f"{field_name}: {field_data['value']} ({field_data.get('confidence', 0):.2f})")
            
            if category_fields:
                print(f"     {category}:")
                for field in category_fields:
                    print(f"       • {field}")


def test_with_roi(client: HFDotsOCRClient, image_path: str) -> None:
    """Test OCR with different ROI regions."""
    print(f"\n🎯 Testing with different ROI regions...")
    
    # Define different ROI regions to test
    roi_regions = {
        "Full Image": None,
        "Top Half": {"x1": 0.0, "y1": 0.0, "x2": 1.0, "y2": 0.5},
        "Bottom Half": {"x1": 0.0, "y1": 0.5, "x2": 1.0, "y2": 1.0},
        "Center Region": {"x1": 0.25, "y1": 0.25, "x2": 0.75, "y2": 0.75},
        "Left Side": {"x1": 0.0, "y1": 0.0, "x2": 0.5, "y2": 1.0},
        "Right Side": {"x1": 0.5, "y1": 0.0, "x2": 1.0, "y2": 1.0}
    }
    
    for region_name, roi in roi_regions.items():
        print(f"\n📍 Testing {region_name}...")
        result = client.extract_text(image_path, roi)
        
        if "error" not in result:
            print(f"   ✅ Success - Processing time: {result.get('processing_time', 0):.2f}s")
            # Show a summary of extracted fields
            detections = result.get('detections', [])
            if detections:
                fields = detections[0].get('extracted_fields', {})
                field_count = sum(1 for field in fields.values() if field and field.get('value'))
                print(f"   📊 Extracted {field_count} fields")
        else:
            print(f"   ❌ Error: {result['error']}")


def main():
    """Main test function."""
    print("🚀 Testing Hugging Face Dots-OCR Space API")
    print("=" * 50)
    
    # Initialize client
    client = HFDotsOCRClient()
    
    # Test 1: Health check
    print("\n1️⃣ Testing health check...")
    health = client.health_check()
    if "error" in health:
        print(f"❌ Health check failed: {health['error']}")
        print("💡 Make sure your Hugging Face Space is running and accessible")
        return
    else:
        print(f"✅ Space is healthy: {health}")
    
    # Test 2: OCR extraction (if demo images exist)
    demo_images = [
        "data/demo/tom_id_card_front.jpg",
        "data/demo/tom_id_card_back.jpg",
        "data/demo/ocr/0000095097_1_E-5858-MA Fahrzeugschein und -brief.png",
        "data/demo/ocr/container_inspection_report.png",
        "data/demo/ocr/handelsregister_b.png"
    ]
    
    test_image = None
    for image_path in demo_images:
        if Path(image_path).exists():
            test_image = image_path
            break
    
    if test_image:
        print(f"\n2️⃣ Testing OCR extraction with {test_image}...")
        result = client.extract_text(test_image)
        print_ocr_results(result)
        
        # Test 3: ROI testing
        if "error" not in result:
            test_with_roi(client, test_image)
    else:
        print("\n2️⃣ No demo images found for testing")
        print("💡 Place some test images in the data/demo/ directory")
    
    print("\n🎉 Testing complete!")
    print("\n💡 To test with your own files:")
    print("   python test_hf_dots_ocr_space.py")
    print("\n💡 To test with ROI:")
    print("   client = HFDotsOCRClient()")
    print("   result = client.extract_text('image.jpg', roi={'x1': 0.0, 'y1': 0.0, 'x2': 0.5, 'y2': 0.5})")


if __name__ == "__main__":
    main()