#!/usr/bin/env python3
"""
OCR toàn bộ Sadi-191-195 sử dụng Qwen VL 3.5 397B
"""

import os
import base64
import requests

# Cấu hình
PDF_NAME = "Sadi-191-195"
WORKSPACE = "/opt/openclaw/.openclaw/workspace/Chuan Muc Sadi"
OUTPUT_FILE = f"{WORKSPACE}/extracted/{PDF_NAME}-qwen.md"
API_KEY = os.environ.get("OPENROUTER_API_KEY")

if not API_KEY:
    print("❌ Thiếu OPENROUTER_API_KEY")
    exit(1)

# Danh sách các trang
pages = []
for i in range(1, 6):
    page_path = f"{WORKSPACE}/sadi_page-{i}.png"
    if os.path.exists(page_path):
        pages.append(page_path)

print(f"📄 Tìm thấy {len(pages)} trang để OCR")
print(f"📝 Model: qwen/qwen3.5-397b-a17b")
print(f"💾 Output: {OUTPUT_FILE}")

# Function để mã hóa ảnh sang base64
def image_to_base64(image_path):
    with open(image_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")

# Prompt cho OCR
PROMPT = """Hãy OCR trang sách Phật giáo Myanmar này. Trích xuất CHÍNH XÁC NGUYÊN VĂN (literal transcription) toàn bộ văn bản.

QUY TẮC BẮT BUỘC:
1. Giữ nguyên mọi ký tự, kể cả lỗi in ấn, ký tự cổ,写法 cổ
2. Không thêm markdown formatting (headers, bold, lists)
3. Đọc từ trên xuống dưới, trái sang phải
4. Output thuần văn bản Myanmar

Bắt đầu bằng "--- TRANG X ---" và kết thúc bằng "--- HẾT TRANG X ---"
"""

# Chuẩn bị headers
headers = {
    "Authorization": f"Bearer {API_KEY}",
    "Content-Type": "application/json",
    "HTTP-Referer": "https://openclaw.ai",
    "X-Title": "OpenClaw OCR Myanmar"
}

# Xử lý từng trang
all_results = []

for idx, page_path in enumerate(pages, 1):
    print(f"\n🔄 [{idx}/{len(pages)}] Đang OCR {os.path.basename(page_path)}...")
    
    # Mã hóa ảnh
    base64_image = image_to_base64(page_path)
    print(f"   📊 Kích thước base64: {len(base64_image)} ký tự")
    
    # Chuẩn bị payload
    payload = {
        "model": "qwen/qwen3.5-397b-a17b",
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": PROMPT.replace("X", str(idx))},
                    {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{base64_image}"}}
                ]
            }
        ],
        "max_tokens": 4096,
        "temperature": 0.1
    }
    
    # Gửi request
    try:
        response = requests.post(
            "https://openrouter.ai/api/v1/chat/completions",
            headers=headers,
            json=payload,
            timeout=180
        )
        
        if response.status_code == 200:
            result = response.json()
            content = result["choices"][0]["message"]["content"]
            all_results.append(content)
            print(f"   ✅ Hoàn thành ({len(content)} ký tự)")
        else:
            error_msg = f"❌ Lỗi API: {response.status_code} - {response.text[:200]}"
            print(f"   {error_msg}")
            all_results.append(f"[LỖI TRANG {idx}: {response.status_code}]")
            
    except Exception as e:
        error_msg = f"❌ Ngoại lệ: {str(e)}"
        print(f"   {error_msg}")
        all_results.append(f"[LỖI TRANG {idx}: {str(e)}]")

# Ghi kết quả ra file
print(f"\n💾 Đang lưu kết quả...")

output_content = f"""# OCR Results - {PDF_NAME}
**Model:** Qwen VL 3.5 397B (openrouter/qwen3.5-397b-a17b)
**Date:** 2026-04-24
**Pages:** {len(pages)}
**Source:** {PDF_NAME}.pdf

---

""" + "\n\n".join(all_results)

with open(OUTPUT_FILE, "w", encoding="utf-8") as f:
    f.write(output_content)

print(f"✅ Hoàn thành! Đã lưu vào {OUTPUT_FILE}")
print(f"📊 Tổng số trang: {len(pages)}")
print(f"📊 Tổng kích thước output: {len(output_content)} ký tự")