#!/usr/bin/env python3
"""
OCR script for Sadi batch - Converts PDF to images and OCRs with OpenRouter Qwen 3.5
"""
import os
import sys
import base64
import requests
import fitz  # PyMuPDF

def pdf_to_images(pdf_path, output_dir):
    """Convert PDF pages to PNG images"""
    os.makedirs(output_dir, exist_ok=True)
    doc = fitz.open(pdf_path)
    images = []
    
    for i, page in enumerate(doc):
        # Render page to image (high resolution)
        mat = fitz.Matrix(2, 2)  # 2x zoom for better OCR
        pix = page.get_pixmap(matrix=mat)
        img_path = os.path.join(output_dir, f"page_{i+1:03d}.png")
        pix.save(img_path)
        images.append(img_path)
        print(f"✅ Converted page {i+1}/{len(doc)}")
    
    doc.close()
    return images

def image_to_base64(img_path):
    """Convert image to base64"""
    with open(img_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")

def ocr_image(base64_image, api_key):
    """OCR single image with OpenRouter Qwen 3.5"""
    url = "https://openrouter.ai/api/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://openclaw.ai",
        "X-Title": "OpenClaw OCR"
    }
    
    payload = {
        "model": "qwen/qwen3.5-397b-a17b",
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "Trích xuất toàn bộ văn bản tiếng Myanmar từ hình ảnh này. YÊU CỰC KỲ QUAN TRỌNG: Phải giữ nguyên văn tuyệt đối (literal transcription) - không sửa lỗi, không làm sạch, giữ nguyên ký tự lặp, ký tự cổ, lỗi chính tả nếu có. Chỉ trả về văn bản Myanmar, không thêm giải thích."
                    },
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/png;base64,{base64_image}"}
                    }
                ]
            }
        ],
        "max_tokens": 4000
    }
    
    response = requests.post(url, headers=headers, json=payload, timeout=120)
    response.raise_for_status()
    result = response.json()
    return result["choices"][0]["message"]["content"]

def main():
    if len(sys.argv) < 2:
        print("Usage: python ocr_sadi_batch.py <pdf_path>")
        sys.exit(1)
    
    pdf_path = sys.argv[1]
    api_key = os.environ.get("OPENROUTER_API_KEY")
    
    if not api_key:
        # Try to load from .env
        env_path = os.path.join(os.path.dirname(os.path.dirname(pdf_path)), ".env")
        if os.path.exists(env_path):
            with open(env_path) as f:
                for line in f:
                    if line.startswith("OPENROUTER_API_KEY="):
                        api_key = line.split("=", 1)[1].strip().strip('"')
                        break
    
    if not api_key:
        print("❌ Error: OPENROUTER_API_KEY not found")
        sys.exit(1)
    
    # Convert PDF to images
    base_name = os.path.splitext(os.path.basename(pdf_path))[0]
    output_dir = f"/tmp/sadi_ocr_{base_name}"
    output_md = os.path.join(os.path.dirname(pdf_path), f"../{base_name}-qwen.md")
    
    print(f"📖 Processing: {pdf_path}")
    print(f"📁 Output dir: {output_dir}")
    
    images = pdf_to_images(pdf_path, output_dir)
    
    # OCR each page
    results = []
    for i, img_path in enumerate(images):
        print(f"🔍 OCR page {i+1}/{len(images)}...")
        base64_img = image_to_base64(img_path)
        text = ocr_image(base64_img, api_key)
        results.append(f"## Trang {i+1}\n\n{text}\n")
        print(f"✅ Page {i+1} done")
    
    # Save results
    os.makedirs(os.path.dirname(output_md), exist_ok=True)
    with open(output_md, "w", encoding="utf-8") as f:
        f.write(f"# {base_name} - OCR Result\n\n")
        f.write("\n".join(results))
    
    print(f"\n✅ OCR complete! Saved to: {output_md}")
    
    # Cleanup
    import shutil
    shutil.rmtree(output_dir)
    print("🧹 Cleaned up temp files")

if __name__ == "__main__":
    main()