"""Example: Batch OCR processing of a directory of images."""

from unlimited_ocr import OCRPipeline

# Initialize the pipeline
pipeline = OCRPipeline(
    model_path="AutomatosX/AX-Unlimited-OCR-3B-MoE-MLX-MXFP8",
    verbose=True,
)

# --- Process all images in a directory ---
# Supports: .jpg, .jpeg, .png, .tiff, .tif, .webp, .bmp
results = pipeline.run_batch(
    "./scanned_documents/",
    format="text",
    show_progress=True,  # Rich progress bar
)

# Results is a dict: {filename: ocr_text}
for filename, text in results.items():
    print(f"\n{'='*60}")
    print(f"FILE: {filename}")
    print(f"{'='*60}")
    print(text[:500])  # Print first 500 chars

# --- Save each result to an output directory ---
results = pipeline.run_batch(
    "./scanned_documents/",
    format="markdown",
    output_dir="./ocr_results/",  # Creates .md files per image
    show_progress=True,
)
print(f"\nProcessed {len(results)} files → ./ocr_results/")

# --- With preprocessing for low-quality scans ---
results = pipeline.run_batch(
    "./low_quality_scans/",
    format="text",
    preprocess=True,  # deskew + contrast enhancement
    output_dir="./cleaned_results/",
)

# --- JSON output with bounding boxes ---
results = pipeline.run_batch(
    "./forms/",
    format="json",
    grounding=True,  # Include bounding box coordinates
    output_dir="./structured_results/",
)

# Clean up
pipeline.cleanup()