#!/usr/bin/env python3
"""
Simple OCR script for Sadi-136-140
"""
import os
import sys
import time
from pathlib import Path

# Add workspace to path
sys.path.insert(0, '/opt/openclaw/.openclaw/workspace/scripts')

from pdf2image import convert_from_path
from google.cloud import vision

BASE_DIR = "/opt/openclaw/.openclaw/workspace/Chuan Muc Sadi"
EXTRACTED_DIR = BASE_DIR + "/extracted"
GOOGLE_SA_KEY = "/opt/openclaw/.openclaw/workspace/google_service_account.json"

def main():
    batch = "136-140"
    pdf_path = os.path.join(BASE_DIR, f"Sadi-{batch}.pdf")
    output_path = os.path.join(EXTRACTED_DIR, f"Sadi-{batch}-qwen.md")
    temp_dir = os.path.join(EXTRACTED_DIR, f"temp_{batch}")
    
    print(f"[START] OCR batch Sadi-{batch}")
    start_time = time.time()
    
    # Check PDF exists
    if not os.path.exists(pdf_path):
        print(f"ERROR: PDF not found: {pdf_path}")
        sys.exit(1)
    
    # Create temp dir
    Path(temp_dir).mkdir(parents=True, exist_ok=True)
    
    # Step 1: Convert PDF to images
    print(f"[STEP 1] Converting PDF to images...")
    try:
        images = convert_from_path(pdf_path, dpi=300)
        print(f"  Converted {len(images)} pages")
        
        image_paths = []
        for i, image in enumerate(images):
            img_path = os.path.join(temp_dir, f"page_{i+1:03d}.png")
            image.save(img_path, 'PNG')
            image_paths.append(img_path)
            print(f"  Saved: {img_path}")
    except Exception as e:
        print(f"ERROR converting PDF: {e}")
        sys.exit(1)
    
    # Step 2: OCR with Google Vision
    print(f"[STEP 2] Running Google Vision OCR...")
    os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = GOOGLE_SA_KEY
    
    try:
        client = vision.ImageAnnotatorClient()
    except Exception as e:
        print(f"ERROR creating Vision client: {e}")
        sys.exit(1)
    
    results = []
    for i, img_path in enumerate(image_paths):
        print(f"  OCR page {i+1}/{len(image_paths)}...")
        
        try:
            with open(img_path, 'rb') as f:
                content = f.read()
            
            image = vision.Image(content=content)
            context = vision.ImageContext(language_hints=['my'])
            response = client.text_detection(image=image, image_context=context)
            
            texts = response.text_annotations
            if texts and len(texts) > 0:
                full_text = texts[0].description
                results.append(f"--- Trang {i+1} ---\n{full_text}")
                print(f"    Extracted {len(full_text)} characters")
            else:
                results.append(f"--- Trang {i+1} ---\n[Không phát hiện văn bản]")
        except Exception as e:
            results.append(f"--- Trang {i+1} ---\n[ERROR: {e}]")
            print(f"    ERROR: {e}")
    
    # Step 3: Save results
    print(f"[STEP 3] Saving results to {output_path}...")
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write('\n\n'.join(results))
    
    # Cleanup
    print(f"[CLEANUP] Removing temp files...")
    import shutil
    try:
        shutil.rmtree(temp_dir)
        print(f"  Removed: {temp_dir}")
    except Exception as e:
        print(f"  Warning: Could not remove temp dir: {e}")
    
    elapsed = time.time() - start_time
    print(f"[DONE] Completed in {elapsed:.1f} seconds")
    print(f"Output: {output_path}")

if __name__ == "__main__":
    main()