#!/usr/bin/env python3
"""
Edit/proofread Myanmar OCR result with OpenRouter Qwen 3.5
"""
import os
import sys
import base64
import requests

def read_file(path):
    with open(path, "r", encoding="utf-8") as f:
        return f.read()

def edit_text(text, api_key):
    """Edit Myanmar text 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 Edit"
    }
    
    prompt = f"""Bạn là chuyên gia hiệu đính văn bản Phật giáo Myanmar.

Nhiệm vụ: HIỆU ĐÍNHH văn bản Myanmar dưới đây. Yêu cầu BẮT BUỘC:

1. GIỮ NGUYÊN 100% cấu trúc, định dạng, số trang, số câu
2. Chỉ sửa lỗi OCR rõ ràng sau:
   - သံဃေန (đúng) ← không phải သံဂူန, သံကူ, သံကွာ
   - ဓမ္မီယ (đúng) ← không phải ဓမ္မိယ
   - ကရဏီယာ (đúng) ← không phải ကရဏ္ဏိယာ, ကရက်ယာ
   - အန္တရေ (đúng) ← không phải အန္ဓရာရေ, အန္တရေဓ
   - နိဗ္ဗိစ္ဆာမိ (đúng) ← không phải နိဗ္ဗိသာမိ, နံသိသာဝမိ
   - ဂမိစ္ဆာမိ (đúng) ← không phải ဂံသွာမိ, ဂံသာမိ, ဂမိသာဝမိ
3. GIỮ NGUYÊN từ Pali, Myanmar cổ, chú thích tiếng Myanmar
4. Trả về TOÀN BỘ văn bản đã hiệu đính, không tóm tắt, không giải thích

Văn bản cần hiệu đính:
{text}"""
    
    payload = {
        "model": "qwen/qwen3.5-397b-a17b",
        "messages": [
            {"role": "user", "content": prompt}
        ],
        "max_tokens": 8000
    }
    
    response = requests.post(url, headers=headers, json=payload, timeout=300)
    response.raise_for_status()
    result = response.json()
    return result["choices"][0]["message"]["content"]

def main():
    if len(sys.argv) < 2:
        print("Usage: python edit_myanmar_ocr.py <md_path>")
        sys.exit(1)
    
    md_path = sys.argv[1]
    api_key = os.environ.get("OPENROUTER_API_KEY")
    
    if not api_key:
        env_path = os.path.join(os.path.dirname(md_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)
    
    print(f"📖 Reading: {md_path}")
    text = read_file(md_path)
    
    print("🔍 Editing with Qwen 3.5...")
    edited = edit_text(text, api_key)
    
    # Save result
    base_name = os.path.splitext(md_path)[0]
    output_path = f"{base_name}-edited.md"
    
    with open(output_path, "w", encoding="utf-8") as f:
        f.write(edited)
    
    print(f"✅ Edit complete! Saved to: {output_path}")

if __name__ == "__main__":
    main()