#!/usr/bin/env python3
"""
OCR Sadi-196-200 sử dụng Qwen VL 3.5 397B qua OpenRouter API
"""

import os
import sys
import base64
import requests

# Flush output ngay lập tức
sys.stdout.reconfigure(line_buffering=True)

# Cấu hình
PDF_NAME = "Sadi-196-200"
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 environment variable", flush=True)
    exit(1)

print(f"🔑 API Key: {API_KEY[:10]}...", flush=True)

# Danh sách các trang
pages = []
for i in range(1, 6):
    page_path = f"{WORKSPACE}/sadi196_page-{i}.png"
    if os.path.exists(page_path):
        pages.append(page_path)
        print(f"📄 Tìm thấy trang {i}: {page_path}", flush=True)
    else:
        print(f"⚠️ Không tìm thấy trang {i}: {page_path}", flush=True)

if not pages:
    print("❌ Không tìm thấy trang nào để OCR", flush=True)
    exit(1)

print(f"📄 Tổng cộng {len(pages)} trang để OCR", flush=True)

# 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
SYSTEM_PROMPT = """Bạn là chuyên gia OCR văn bản Myanmar (Burmese) cổ từ sách Phật giáo Theravada.

NHIỆM VỤ: Trích xuất CHÍNH XÁC NGUYÊN VĂN (literal transcription) toàn bộ văn bản Myanmar từ ảnh chụp trang sách.

QUY TẮC BẮT BUỘC:
1. KHÔNG tự ý chỉnh sửa, làm sạch hay thay đổi bất kỳ ký tự nào
2. GIỮ NGUYÊN lỗi lặp từ, ký tự cổ - kể cả khi thấy "sai"
3. KHÔNG thêm markdown formatting (headers, bold, lists) - chỉ output text thuần
4. XUỐNG DÒNG đúng theo bố cục trang sách gốc
5. ĐỌC TỪ TRÊN XUỐNG DƯỚI, từ TRÁI SANG PHẢI

LƯU Ý ĐẶC BIỆT:
- Văn bản Phật pháp Pali-Myanmar có nhiều ký tự cổ, viết tắt
- Có thể có lỗi in ấn từ bản gốc - GIỮ NGUYÊN
- Không thêm ghi chú, comment, hay giải thích
- Output thuần văn bản Myanmar, không dịch, không romanize

ĐỊNH DẠNG OUTPUT:
- Bắt đầu bằng "--- TRANG X ---" (X là số trang)
- Sau đó là toàn bộ văn bản trích xuất được
- Kết thúc bằng "--- HẾT TRANG X ---"
"""

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

for idx, page_path in enumerate(pages, 1):
    print(f"\n🔄 Đang OCR trang {idx}/{len(pages)}: {os.path.basename(page_path)}", flush=True)
    
    # Mã hóa ảnh
    base64_image = image_to_base64(page_path)
    print(f"  📊 Ảnh size: {len(base64_image)} bytes (base64)", flush=True)
    
    # Chuẩn bị request
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://openclaw.ai",
        "X-Title": "OpenClaw OCR Myanmar"
    }
    
    payload = {
        "model": "qwen/qwen3.5-397b-a17b",
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": f"Hãy OCR trang {idx} này. Trích xuất chính xác nguyên văn toàn bộ văn bản Myanmar."
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/png;base64,{base64_image}"
                        }
                    }
                ]
            }
        ],
        "max_tokens": 4096,
        "temperature": 0.1
    }
    
    # Gửi request
    try:
        print(f"  🌐 Đang gửi request đến OpenRouter...", flush=True)
        response = requests.post(
            "https://openrouter.ai/api/v1/chat/completions",
            headers=headers,
            json=payload,
            timeout=180
        )
        
        print(f"  📡 Response status: {response.status_code}", flush=True)
        
        if response.status_code == 200:
            result = response.json()
            content = result["choices"][0]["message"]["content"]
            all_results.append(content)
            print(f"  ✅ Trang {idx} hoàn thành ({len(content)} ký tự)", flush=True)
        else:
            error_msg = f"❌ Lỗi API trang {idx}: {response.status_code} - {response.text[:200]}"
            print(error_msg, flush=True)
            all_results.append(f"[LỖI TRANG {idx}: {response.status_code}]")
            
    except Exception as e:
        error_msg = f"❌ Ngoại lệ trang {idx}: {str(e)}"
        print(error_msg, flush=True)
        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ả vào {OUTPUT_FILE}...", flush=True)

output_content = f"""# OCR Results - {PDF_NAME}
**Model:** Qwen VL 3.5 397B (openrouter/qwen-3.5-397b-a22b-instruct)
**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}", flush=True)
print(f"📊 Tổng số trang: {len(pages)}", flush=True)
print(f"📊 Tổng kích thước output: {len(output_content)} ký tự", flush=True)