#!/usr/bin/env python3
"""
Sample client for GLM-OCR API
Demonstrates various usage patterns
"""

import requests
import json
import sys
from pathlib import Path
from typing import Dict, List, Optional

class OCRClient:
    """Client for GLM-OCR API"""
    
    def __init__(self, base_url: str = "http://localhost:7860"):
        self.base_url = base_url
        self.session = requests.Session()
    
    def health_check(self) -> Dict:
        """Check if API is running"""
        try:
            response = self.session.get(f"{self.base_url}/health")
            response.raise_for_status()
            return response.json()
        except requests.RequestException as e:
            print(f"❌ Health check failed: {e}")
            return None
    
    def get_info(self) -> Dict:
        """Get API information"""
        try:
            response = self.session.get(f"{self.base_url}/api/info")
            response.raise_for_status()
            return response.json()
        except requests.RequestException as e:
            print(f"❌ Failed to get info: {e}")
            return None
    
    def extract_text(self, file_path: str) -> Optional[Dict]:
        """
        Extract text, tables, and formulas from document.
        
        Args:
            file_path: Path to document file
        
        Returns:
            Extraction results or None if failed
        """
        if not Path(file_path).exists():
            print(f"❌ File not found: {file_path}")
            return None
        
        try:
            with open(file_path, 'rb') as f:
                files = {'file': f}
                response = self.session.post(
                    f"{self.base_url}/api/ocr/extract",
                    files=files,
                    timeout=120
                )
            
            response.raise_for_status()
            return response.json()
        
        except requests.RequestException as e:
            print(f"❌ Extraction failed: {e}")
            return None
    
    def extract_structured(
        self,
        file_path: str,
        schema: Dict
    ) -> Optional[Dict]:
        """
        Extract structured information using JSON schema.
        
        Args:
            file_path: Path to document file
            schema: JSON schema defining extraction structure
        
        Returns:
            Extraction results or None if failed
        """
        if not Path(file_path).exists():
            print(f"❌ File not found: {file_path}")
            return None
        
        try:
            with open(file_path, 'rb') as f:
                files = {'file': f}
                data = {'schema': json.dumps(schema)}
                
                response = self.session.post(
                    f"{self.base_url}/api/ocr/extract-structured",
                    files=files,
                    data=data,
                    timeout=120
                )
            
            response.raise_for_status()
            return response.json()
        
        except requests.RequestException as e:
            print(f"❌ Structured extraction failed: {e}")
            return None
    
    def batch_extract(self, file_paths: List[str]) -> Optional[Dict]:
        """
        Process multiple files in batch.
        
        Args:
            file_paths: List of file paths to process
        
        Returns:
            Batch results or None if failed
        """
        if len(file_paths) > 10:
            print("❌ Maximum 10 files per batch")
            return None
        
        files = []
        for file_path in file_paths:
            if not Path(file_path).exists():
                print(f"⚠️  File not found: {file_path}")
                continue
            
            try:
                files.append(('files', open(file_path, 'rb')))
            except IOError as e:
                print(f"⚠️  Cannot open {file_path}: {e}")
        
        if not files:
            print("❌ No valid files to process")
            return None
        
        try:
            response = self.session.post(
                f"{self.base_url}/api/ocr/batch",
                files=files,
                timeout=300
            )
            
            # Close all file handles
            for _, file_obj in files:
                file_obj.close()
            
            response.raise_for_status()
            return response.json()
        
        except requests.RequestException as e:
            print(f"❌ Batch extraction failed: {e}")
            return None


def print_extraction_results(results: Dict):
    """Pretty print extraction results"""
    print("\n" + "="*80)
    print(f"📄 File: {results['filename']}")
    print(f"📋 Type: {results['file_type']}")
    print(f"📑 Pages: {results['total_pages']}")
    print("="*80)
    
    for page in results['pages']:
        print(f"\n📖 Page {page['page_number']}:")
        print("-" * 80)
        
        if page.get('table'):
            print("\n🔷 TABLE CONTENT:")
            print(page['table'][:500] + ("..." if len(page['table']) > 500 else ""))
        
        if page.get('text'):
            print("\n📝 TEXT CONTENT:")
            print(page['text'][:500] + ("..." if len(page['text']) > 500 else ""))
        
        if page.get('formula'):
            print("\n📐 FORMULA CONTENT:")
            print(page['formula'][:500] + ("..." if len(page['formula']) > 500 else ""))


def print_structured_results(results: Dict):
    """Pretty print structured extraction results"""
    print("\n" + "="*80)
    print(f"📄 File: {results['filename']}")
    print(f"📋 Type: {results['file_type']}")
    print("="*80)
    
    for page in results['pages']:
        print(f"\n📖 Page {page['page_number']}:")
        print("-" * 80)
        
        extracted = page.get('extracted_data', {})
        print(json.dumps(extracted, indent=2, ensure_ascii=False))


def main():
    """Main example function"""
    
    # Initialize client
    client = OCRClient()
    
    # Check API health
    print("🔍 Checking API health...")
    health = client.health_check()
    if health:
        print(f"✅ API is healthy: {health['status']}")
    else:
        print("❌ Cannot connect to API. Make sure it's running on http://localhost:7860")
        sys.exit(1)
    
    # Get API info
    print("\n📚 Getting API information...")
    info = client.get_info()
    if info:
        print(f"✅ API: {info['api_name']} v{info['version']}")
        print(f"   Model: {info['model']}")
        print(f"   Supported formats: {', '.join(info['supported_formats']['images'] + info['supported_formats']['documents'] + info['supported_formats']['pdfs'])}")
    
    # Example 1: Extract from a single image/document
    print("\n" + "="*80)
    print("EXAMPLE 1: Single File Extraction")
    print("="*80)
    
    # Replace with your actual file path
    sample_file = "sample_document.pdf"
    
    if Path(sample_file).exists():
        print(f"\n📤 Extracting from: {sample_file}")
        results = client.extract_text(sample_file)
        
        if results:
            print_extraction_results(results)
            
            # Save results to file
            output_file = Path(sample_file).stem + "_extracted.json"
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(results, f, indent=2, ensure_ascii=False)
            print(f"\n💾 Results saved to: {output_file}")
    else:
        print(f"\n⚠️  Sample file not found: {sample_file}")
        print("   Create a test file or provide a file path")
    
    # Example 2: Structured extraction
    print("\n" + "="*80)
    print("EXAMPLE 2: Structured Information Extraction")
    print("="*80)
    
    # Define extraction schema for ID card
    id_schema = {
        "id_number": "",
        "full_name": "",
        "date_of_birth": "",
        "address": {
            "street": "",
            "city": "",
            "state": "",
            "zip_code": ""
        },
        "issue_date": "",
        "expiration_date": ""
    }
    
    # Replace with your ID document path
    id_file = "id_card.jpg"
    
    if Path(id_file).exists():
        print(f"\n📤 Extracting structured data from: {id_file}")
        results = client.extract_structured(id_file, id_schema)
        
        if results:
            print_structured_results(results)
            
            # Save results
            output_file = Path(id_file).stem + "_structured.json"
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(results, f, indent=2, ensure_ascii=False)
            print(f"\n💾 Results saved to: {output_file}")
    else:
        print(f"\n⚠️  ID file not found: {id_file}")
        print("   For structured extraction, provide a relevant document")
    
    # Example 3: Batch processing
    print("\n" + "="*80)
    print("EXAMPLE 3: Batch Processing")
    print("="*80)
    
    # List of files to process
    files_to_batch = [
        "document1.pdf",
        "document2.jpg",
        "document3.docx"
    ]
    
    # Filter existing files
    existing_files = [f for f in files_to_batch if Path(f).exists()]
    
    if existing_files:
        print(f"\n📤 Processing batch: {len(existing_files)} files")
        results = client.batch_extract(existing_files)
        
        if results:
            print(f"✅ Batch timestamp: {results['batch_timestamp']}")
            print(f"   Total files: {results['total_files']}")
            
            for file_result in results['results']:
                if file_result['status'] == 'success':
                    print(f"   ✅ {file_result['filename']} - {file_result['total_pages']} pages")
                else:
                    print(f"   ❌ {file_result['filename']} - {file_result.get('error', 'Unknown error')}")
            
            # Save batch results
            with open("batch_results.json", 'w', encoding='utf-8') as f:
                json.dump(results, f, indent=2, ensure_ascii=False)
            print(f"\n💾 Results saved to: batch_results.json")
    else:
        print(f"\n⚠️  No files found for batch processing")
        print("   Prepare test documents and update file paths in this script")
    
    print("\n" + "="*80)
    print("✅ Examples complete!")
    print("="*80)


if __name__ == "__main__":
    main()