#!/usr/bin/env python3
"""
OCR PDF to RAW JSON using Google Cloud Vision Async API.
Usage: python3 ocr_pdf_to_raw.py <local_pdf_path> [bucket_name] [local_output_dir]
"""
import os
import sys
import time
import json
from google.cloud import vision
from google.cloud import storage

# === CONFIG ===
KEY_FILE = '/home/tuan-nguyen/.openclaw/workspace/old/google_service_account.json'
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = KEY_FILE

BUCKET_NAME = 'zen-ocr-pdf'
BATCH_SIZE = 3  # pages per JSON output file (reduced from 5 to avoid editor overload)
TIMEOUT = 900    # seconds (15 min)

def upload_pdf(local_path, bucket_name):
    """Upload PDF to GCS, return GCS URI."""
    client = storage.Client()
    bucket = client.bucket(bucket_name)
    pdf_name = os.path.basename(local_path)
    blob = bucket.blob(f"input/{pdf_name}")
    
    print(f"⬆️  Uploading {local_path} → gs://{bucket_name}/input/{pdf_name} ...")
    blob.upload_from_filename(local_path)
    
    uri = f"gs://{bucket_name}/input/{pdf_name}"
    print(f"✅ Uploaded: {uri} ({blob.size:,} bytes)")
    return uri


def ocr_async(gcs_source_uri, gcs_dest_uri):
    """Call Cloud Vision asyncBatchAnnotateFiles."""
    client = vision.ImageAnnotatorClient()

    feature = vision.Feature(type_=vision.Feature.Type.DOCUMENT_TEXT_DETECTION)
    
    input_config = vision.InputConfig(
        gcs_source=vision.GcsSource(uri=gcs_source_uri),
        mime_type='application/pdf'
    )
    
    output_config = vision.OutputConfig(
        gcs_destination=vision.GcsDestination(uri=gcs_dest_uri),
        batch_size=BATCH_SIZE
    )
    
    image_context = vision.ImageContext(
        language_hints=['my', 'pi', 'en']
    )
    
    async_request = vision.AsyncAnnotateFileRequest(
        features=[feature],
        input_config=input_config,
        output_config=output_config,
        image_context=image_context,
    )
    
    print(f"🚀 Starting async OCR...")
    print(f"   Input:  {gcs_source_uri}")
    print(f"   Output: {gcs_dest_uri}")
    print(f"   Batch size: {BATCH_SIZE} pages/file")
    
    operation = client.async_batch_annotate_files(requests=[async_request])
    print(f"   Operation: {operation.operation.name}")
    print(f"⏳ Waiting for completion (timeout={TIMEOUT}s)...")
    
    start = time.time()
    result = operation.result(timeout=TIMEOUT)
    elapsed = time.time() - start
    
    print(f"✅ OCR completed in {elapsed:.0f}s")
    return result


def download_results(gcs_output_prefix, local_dir):
    """Download all JSON output files from GCS to local. NO merging."""
    client = storage.Client()
    
    parts = gcs_output_prefix.replace('gs://', '').split('/', 1)
    bucket_name = parts[0]
    prefix = parts[1] if len(parts) > 1 else ''
    
    bucket = client.bucket(bucket_name)
    blobs = list(bucket.list_blobs(prefix=prefix))
    json_blobs = [b for b in blobs if b.name.endswith('.json')]
    
    if not json_blobs:
        print("⚠️  No JSON output files found!")
        return
    
    os.makedirs(local_dir, exist_ok=True)
    
    total_bytes = 0
    for blob in json_blobs:
        filename = os.path.basename(blob.name)
        local_path = os.path.join(local_dir, filename)
        blob.download_to_filename(local_path)
        total_bytes += blob.size
        print(f"   ✅ {filename} ({blob.size:,} bytes)")
    
    print(f"📥 Downloaded {len(json_blobs)} JSON files ({total_bytes:,} bytes total) → {local_dir}")


def main():
    if len(sys.argv) < 2:
        print("Usage: python3 ocr_pdf_to_raw.py <local_pdf_path> [bucket_name] [local_output_dir]")
        sys.exit(1)
    
    local_pdf_path = sys.argv[1]
    bucket_name = sys.argv[2] if len(sys.argv) > 2 else BUCKET_NAME
    local_output_dir = sys.argv[3] if len(sys.argv) > 3 else 'ocr/raw/'
    
    if not os.path.exists(local_pdf_path):
        print(f"❌ File not found: {local_pdf_path}")
        sys.exit(1)
    
    pdf_name = os.path.splitext(os.path.basename(local_pdf_path))[0]
    
    print("=" * 60)
    print(f"📄 OCR Pipeline: {os.path.basename(local_pdf_path)}")
    print(f"   Batch size: {BATCH_SIZE} pages/file")
    print(f"   Bucket:     gs://{bucket_name}")
    print(f"   Output:     {os.path.abspath(local_output_dir)}")
    print("=" * 60)
    
    # Step 1: Upload
    gcs_source_uri = upload_pdf(local_pdf_path, bucket_name)
    
    # Step 2: OCR
    gcs_dest_uri = f"gs://{bucket_name}/output/{pdf_name}/"
    result = ocr_async(gcs_source_uri, gcs_dest_uri)
    
    # Step 3: Download
    download_results(gcs_dest_uri, local_output_dir)
    
    print(f"\n🎉 Done! JSON files saved to {os.path.abspath(local_output_dir)}")


if __name__ == '__main__':
    main()