#!/usr/bin/env python3
"""
OCR Layout Detection API Client
================================
Simple script to interact with the OCR Layout Detection service.

Usage:
    python api_client.py <path_to_file>
    
Examples:
    python api_client.py invoice.pdf
    python api_client.py document.jpg
    python api_client.py signature.png --signature-only
"""

import os
import sys
import json
import argparse
from pathlib import Path
from gradio_client import Client, handle_file


# API Configuration
SPACE_URL = "Ayaan-Sharif/ocr-layout-detection-poc"
HF_TOKEN = os.environ.get("HF_TOKEN")  # Read from environment variable if available


def analyze_document(file_path, mode="Fast", enable_ocr=True, enable_tables=True, 
                     detect_signatures=False, signature_conf=0.05):
    """
    Analyze a document with layout detection and optional OCR.
    
    Args:
        file_path: Path to PDF or image file
        mode: "Fast" or "Accurate" processing mode
        enable_ocr: Extract text with OCR
        enable_tables: Detect and extract tables
        detect_signatures: Also detect signatures (slower)
        signature_conf: Confidence threshold for signatures (0.01-0.5)
        
    Returns:
        dict: Contains visualization, summary, markdown, and JSON outputs
    """
    print(f"📄 Analyzing document: {file_path}")
    print(f"   Mode: {mode} | OCR: {enable_ocr} | Tables: {enable_tables} | Signatures: {detect_signatures}")
    
    try:
        client = Client(SPACE_URL, hf_token=HF_TOKEN)
        
        result = client.predict(
            file=handle_file(file_path),
            mode=mode,
            enable_ocr=enable_ocr,
            enable_tables=enable_tables,
            run_signature_yolo=detect_signatures,
            signature_conf=signature_conf,
            api_name="/gradio_interface"
        )
        
        # result is a tuple: (visualization_image, summary_text, markdown_text, json_text)
        visualization, summary, markdown, json_output = result
        
        print("✅ Analysis complete!")
        return {
            "visualization": visualization,
            "summary": summary,
            "markdown": markdown,
            "json": json_output
        }
        
    except Exception as e:
        print(f"❌ Error: {e}")
        return None


def detect_signatures_only(file_path, multiscale=True, conf=0.03, iou=0.45, augment=True):
    """
    Detect signatures only (faster, no OCR or layout analysis).
    
    Args:
        file_path: Path to PDF or image file
        multiscale: Try multiple scales (1.0, 1.5, 2.0) for better detection
        conf: Confidence threshold (0.01-0.5, lower = more detections)
        iou: IoU threshold for NMS (0.1-0.9)
        augment: Use augmentation (slower but better recall)
        
    Returns:
        dict: Contains annotated image, summary, and JSON detections
    """
    print(f"✍️ Detecting signatures in: {file_path}")
    print(f"   Multiscale: {multiscale} | Conf: {conf} | IoU: {iou} | Augment: {augment}")
    
    try:
        client = Client(SPACE_URL, hf_token=HF_TOKEN)
        
        result = client.predict(
            file=handle_file(file_path),
            try_scales=multiscale,
            conf=conf,
            iou=iou,
            augment=augment,
            api_name="/signature_only_infer"
        )
        
        # result is a tuple: (annotated_image, summary_text, json_detections)
        annotated_image, summary, json_output = result
        
        print("✅ Signature detection complete!")
        return {
            "annotated_image": annotated_image,
            "summary": summary,
            "json": json_output
        }
        
    except Exception as e:
        print(f"❌ Error: {e}")
        return None


def save_results(results, output_dir="output"):
    """Save API results to files."""
    os.makedirs(output_dir, exist_ok=True)
    
    if results is None:
        return
    
    # Save visualization/annotated image
    if "visualization" in results and results["visualization"]:
        viz_path = results["visualization"].get("path")
        if viz_path and os.path.exists(viz_path):
            import shutil
            output_path = os.path.join(output_dir, "visualization.png")
            shutil.copy(viz_path, output_path)
            print(f"💾 Saved visualization: {output_path}")
    
    if "annotated_image" in results and results["annotated_image"]:
        img_path = results["annotated_image"].get("path")
        if img_path and os.path.exists(img_path):
            import shutil
            output_path = os.path.join(output_dir, "signatures_annotated.png")
            shutil.copy(img_path, output_path)
            print(f"💾 Saved annotated image: {output_path}")
    
    # Save markdown content
    if "markdown" in results and results["markdown"]:
        markdown_path = os.path.join(output_dir, "content.md")
        with open(markdown_path, "w", encoding="utf-8") as f:
            f.write(results["markdown"])
        print(f"💾 Saved markdown: {markdown_path}")
    
    # Save JSON output
    if "json" in results and results["json"]:
        json_path = os.path.join(output_dir, "layout.json")
        with open(json_path, "w", encoding="utf-8") as f:
            f.write(results["json"])
        print(f"💾 Saved JSON: {json_path}")
    
    # Save summary
    if "summary" in results and results["summary"]:
        summary_path = os.path.join(output_dir, "summary.txt")
        with open(summary_path, "w", encoding="utf-8") as f:
            f.write(results["summary"])
        print(f"💾 Saved summary: {summary_path}")


def main():
    parser = argparse.ArgumentParser(
        description="OCR Layout Detection API Client",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Full document analysis with OCR
  python api_client.py invoice.pdf
  
  # Accurate mode with signature detection
  python api_client.py document.pdf --mode Accurate --detect-signatures
  
  # Signature detection only (faster)
  python api_client.py contract.jpg --signature-only
  
  # Custom output directory
  python api_client.py file.pdf --output results/
        """
    )
    
    parser.add_argument("file", help="Path to document (PDF, JPG, PNG)")
    parser.add_argument("--mode", choices=["Fast", "Accurate"], default="Fast",
                        help="Processing mode (default: Fast)")
    parser.add_argument("--no-ocr", action="store_true", help="Disable OCR")
    parser.add_argument("--no-tables", action="store_true", help="Disable table detection")
    parser.add_argument("--detect-signatures", action="store_true",
                        help="Also detect signatures in full analysis")
    parser.add_argument("--signature-conf", type=float, default=0.05,
                        help="Signature confidence threshold (default: 0.05)")
    parser.add_argument("--signature-only", action="store_true",
                        help="Only detect signatures (faster, no OCR)")
    parser.add_argument("--output", "-o", default="output",
                        help="Output directory (default: output)")
    
    args = parser.parse_args()
    
    # Validate file exists
    if not os.path.exists(args.file):
        print(f"❌ Error: File not found: {args.file}")
        sys.exit(1)
    
    # Check file type
    ext = Path(args.file).suffix.lower()
    if ext not in [".pdf", ".jpg", ".jpeg", ".png", ".tiff", ".bmp"]:
        print(f"⚠️  Warning: Unsupported file type: {ext}")
        print("   Supported: .pdf, .jpg, .jpeg, .png, .tiff, .bmp")
    
    print(f"\n🚀 Starting API call to {SPACE_URL}\n")
    
    # Call appropriate API endpoint
    if args.signature_only:
        results = detect_signatures_only(args.file)
    else:
        results = analyze_document(
            args.file,
            mode=args.mode,
            enable_ocr=not args.no_ocr,
            enable_tables=not args.no_tables,
            detect_signatures=args.detect_signatures,
            signature_conf=args.signature_conf
        )
    
    # Save results
    if results:
        print(f"\n📁 Saving results to: {args.output}/")
        save_results(results, args.output)
        print("\n✨ Done!")
    else:
        print("\n❌ Failed to process document")
        sys.exit(1)


if __name__ == "__main__":
    main()