#!/usr/bin/env python3
"""
OCR Myanmar text from PDF page images using Qwen VL via OpenRouter.
Sends images as base64 data URIs, writes payload to temp file to avoid arg limit.
"""
import subprocess, json, os, sys, base64, tempfile, time

API_KEY = os.environ['OPENROUTER_API_KEY']

def get_image_b64_data_uri(path):
    with open(path, 'rb') as f:
        b64 = base64.b64encode(f.read()).decode()
    return f"data:image/jpeg;base64,{b64}"

def call_openrouter(image_path, page_label):
    data_uri = get_image_b64_data_uri(image_path)
    
    payload = {
        "model": "qwen/qwen3.5-397b-a17b",
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": f"Extract all Myanmar text from this image exactly as written. Literal transcription only. Do not clean, correct, or alter any text, characters, or punctuation. Include all spacing and punctuation marks such as ၊ ။ etc. This is page {page_label}."
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": data_uri
                        }
                    }
                ]
            }
        ],
        "max_tokens": 4096
    }
    
    # Write to temp file to avoid "Argument list too long"
    with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
        json.dump(payload, f)
        payload_file = f.name
    
    try:
        cmd = [
            "curl", "-s", "-X", "POST",
            "https://openrouter.ai/api/v1/chat/completions",
            "-H", f"Authorization: Bearer {API_KEY}",
            "-H", "Content-Type: application/json",
            "-H", "HTTP-Referer: https://openclaw.ai",
            "-H", "X-Title: OpenClaw OCR",
            "-d", f"@{payload_file}"
        ]
        
        print(f"Calling API for page {page_label}...", flush=True)
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
        
        if result.returncode != 0:
            print(f"CURL ERROR: {result.stderr}", flush=True)
            return None
        
        try:
            resp = json.loads(result.stdout)
        except json.JSONDecodeError as e:
            print(f"JSON PARSE ERROR: {e}", flush=True)
            print(f"RAW: {result.stdout[:500]}", flush=True)
            return None
        
        if 'choices' not in resp or len(resp['choices']) == 0:
            # Check for error field
            if 'error' in resp:
                print(f"API ERROR: {json.dumps(resp['error'], indent=2)[:500]}", flush=True)
            else:
                print(f"UNEXPECTED RESPONSE: {json.dumps(resp, indent=2)[:500]}", flush=True)
            return None
        
        content = resp['choices'][0]['message']['content']
        return content
    finally:
        os.unlink(payload_file)


if len(sys.argv) < 4:
    print("Usage: python3 ocr_batch.py <image_path> <page_label> <output_md>")
    sys.exit(1)

image_path = sys.argv[1]
page_label = sys.argv[2]
output_md = sys.argv[3]

content = call_openrouter(image_path, page_label)
if content:
    print(f"Result OK, length={len(content)}", flush=True)
    print(f"Preview: {content[:300]}", flush=True)
    
    with open(output_md, 'a') as f:
        f.write(f"### Trang {page_label}\n\n")
        f.write(content.strip())
        f.write("\n\n")
else:
    print(f"FAILED for page {page_label}", flush=True)
    with open(output_md, 'a') as f:
        f.write(f"### Trang {page_label}\n\n")
        f.write("*Lỗi: API không trả kết quả*\n\n")