#!/usr/bin/env python3
"""
Hiệu đính văn bản Myanmar OCR với Qwen 3.5 - Theo trang
"""

import sys
import os
import requests

# Setup paths
workspace = "/home/tuan-nguyen/.openclaw/workspace"
input_path = f"{workspace}/Chuan Muc Sadi/Sadi-206-210-qwen.md"
output_path = f"{workspace}/Chuan Muc Sadi/Sadi-206-210-edited.md"

# API Key
OPENROUTER_API_KEY = "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"

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

def split_into_pages(text):
    """Tách văn bản thành từng trang"""
    pages = []
    current_page = []
    current_page_num = None
    
    for line in text.split('\n'):
        if line.startswith('--- Trang'):
            if current_page and current_page_num is not None:
                pages.append((current_page_num, '\n'.join(current_page)))
            current_page = []
            current_page_num = line.replace('--- Trang ', '').replace(' ---', '').strip()
        elif current_page_num is not None:
            current_page.append(line)
    
    if current_page and current_page_num is not None:
        pages.append((current_page_num, '\n'.join(current_page)))
    
    return pages

def edit_page_with_qwen(page_text, page_num):
    """Hiệu đính một trang với Qwen 3.5"""
    
    headers = {
        'Authorization': f'Bearer {OPENROUTER_API_KEY}',
        'Content-Type': 'application/json',
        'HTTP-Referer': 'https://openclaw.ai',
        'X-Title': 'OpenClaw OCR Editor'
    }
    
    prompt = f"""Bạn là chuyên gia hiệu đính văn bản Phật pháp Myanmar.

Nhiệm vụ: Hiệu đính văn bản OCR trang {page_num} sau đây.

Sửa các lỗi thường gặp:
- Số ၈ bị nhận diện thành ဂ → sửa thành ၈
- လည်းကောင်း → လည်း (rút gọn)
- တောလည်းကောင်းင်း → တောင့်
- မှူး → မှဲ့
- ဗောဓိသျား → ဗောဓိသတ်
- ဂံလာန → ဂိလာန
- သင်္ခေသာ → သင်္ဃာ
- ပဋိသာ၀ီ → ပဋိသံဝီ
- ဥုံသ → ဥဏ္ဌ
- မဟက် → မီး
- ခြင် → ခြင်္
- မေ → မေဃ
- အေနအကံ → အေးအကံ
- မိးပကံ → မီးပကံ

Yêu cầu:
- Giữ nguyên thuật ngữ Pali
- Giữ nguyên số trang
- Chỉ sửa lỗi chính tả Myanmar
- Xuất ra văn bản đã hiệu đính, không thêm giải thích

Văn bản cần hiệu đính:
{page_text}

Văn bản đã hiệu đính:"""

    payload = {
        'model': 'qwen/qwen3.5-397b-a17b',
        'messages': [
            {
                'role': 'user',
                'content': prompt
            }
        ],
        'max_tokens': 4096
    }
    
    response = requests.post(
        'https://openrouter.ai/api/v1/chat/completions',
        headers=headers,
        json=payload
    )
    
    if response.status_code == 200:
        data = response.json()
        edited_text = data['choices'][0]['message']['content']
        return edited_text.strip()
    else:
        print(f"  Lỗi API trang {page_num}: {response.status_code}")
        return f"[Lỗi hiệu đính trang {page_num}]"

def main():
    print(f"🔍 Bắt đầu hiệu đính Sadi-206-210 (theo trang)")
    
    if not os.path.exists(input_path):
        print(f"❌ File không tồn tại: {input_path}")
        sys.exit(1)
    
    # Đọc file input
    print(f"📄 Đọc file input...")
    text = read_file(input_path)
    print(f"  Đã đọc {len(text)} ký tự")
    
    # Tách thành từng trang
    print(f"✂️ Tách thành từng trang...")
    pages = split_into_pages(text)
    print(f"  Tìm thấy {len(pages)} trang")
    
    # Hiệu đính từng trang
    print(f"🤖 Hiệu đính với Qwen 3.5...")
    edited_pages = []
    
    for page_num, page_text in pages:
        print(f"  Hiệu đính trang {page_num}...")
        edited_text = edit_page_with_qwen(page_text, page_num)
        edited_pages.append(f"--- Trang {page_num} ---\n{edited_text}")
    
    # Ghi kết quả
    final_text = '\n\n'.join(edited_pages)
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(final_text)
    
    print(f"\n✅ Hiệu đính hoàn thành! Đã lưu vào: {output_path}")
    print(f"Số ký tự: {len(final_text)}")

if __name__ == "__main__":
    main()