#!/usr/bin/env python3
import subprocess, json, os, sys, base64

API_KEY = os.environ['OPENROUTER_API_KEY']

def encode_pdf(path):
    with open(path, 'rb') as f:
        return base64.b64encode(f.read()).decode()

def call_openrouter(b64_data, page_info_text):
    payload = {
        "model": "qwen/qwen3.5-397b-a17b",
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": f"Extract all Myanmar text from this PDF page exactly as written. Literal transcription only. Do not clean or correct any text, characters, or punctuation. Include all spacing. This is page {page_info_text}."
                    },
                    {
                        "type": "file",
                        "file": {
                            "data": b64_data,
                            "mimeType": "application/pdf"
                        }
                    }
                ]
            }
        ]
    }
    
    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", json.dumps(payload)
    ]
    
    print(f"Calling API for {page_info_text}...", 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:
        print(f"API ERROR: {json.dumps(resp, indent=2)[:500]}", flush=True)
        return None
    
    content = resp['choices'][0]['message']['content']
    return content

# Single page
if len(sys.argv) < 4:
    print("Usage: python3 ocr_batch1.py <pdf_path> <page_label> <output_md>")
    sys.exit(1)

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

b64 = encode_pdf(pdf_path)
print(f"PDF encoded: {len(b64)} chars", flush=True)

content = call_openrouter(b64, 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")