#!/usr/bin/env python3
"""
OCR với Qwen Vision qua OpenRouter API
"""

import sys
import os
import base64
import requests
from pathlib import Path
from pdf2image import convert_from_path

OPENROUTER_API_KEY = os.environ.get('OPENROUTER_API_KEY')
if not OPENROUTER_API_KEY:
    print("❌ Thiếu OPENROUTER_API_KEY environment variable")
    sys.exit(1)

def encode_image_to_base64(image_path):
    with open(image_path, 'rb') as f:
        return base64.b64encode(f.read()).decode('utf-8')

def ocr_with_qwen_vision(image_paths, output_path):
    """OCR với Qwen Vision qua OpenRouter API"""
    
    headers = {
        'Authorization': f'Bearer {OPENROUTER_API_KEY}',
        'Content-Type': 'application/json',
        'HTTP-Referer': 'https://openclaw.ai',
        'X-Title': 'OpenClaw OCR'
    }
    
    results = []
    
    for i, img_path in enumerate(image_paths):
        print(f"  OCR trang {i+1}/{len(image_paths)}...")
        
        base64_image = encode_image_to_base64(img_path)
        
        payload = {
            'model': 'qwen/qwen3.5-397b-a17b',
            'messages': [
                {
                    'role': 'user',
                    'content': [
                        {
                            'type': 'image_url',
                            'image_url': {'url': f'data:image/png;base64,{base64_image}'}
                        },
                        {
                            'type': 'text',
                            'text': '''Trích xuất TOÀN BỘ văn bản Myanmar từ hình ảnh này.

Yêu cầu BẮT BUỘC:
- Literal transcription - giữ nguyên văn tuyệt đối, KHÔNG làm sạch, KHÔNG chuẩn hóa
- Giữ nguyên cả lỗi lặp từ, ký tự cổ, ký tự đặc biệt
- Không thêm giải thích, không thêm nhận xét
- Chỉ xuất văn bản thuần túy'''
                        }
                    ]
                }
            ],
            '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()
            text = data['choices'][0]['message']['content']
            # Loại bỏ markdown code blocks nếu có
            if text.startswith('```'):
                lines = text.split('\n')
                if lines[0].startswith('```'):
                    lines = lines[1:]
                if lines and lines[-1].strip() == '```':
                    lines = lines[:-1]
                text = '\n'.join(lines)
            results.append(f"--- Trang {i+1} ---\n{text}")
        else:
            print(f"    Lỗi API: {response.status_code} - {response.text}")
            results.append(f"--- Trang {i+1} ---\n[Lỗi OCR: {response.status_code}]")
    
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write('\n\n'.join(results))
    
    return True

def main():
    if len(sys.argv) < 2:
        print("Usage: python3 ocr_qwen_vision.py <batch>")
        print("Example: python3 ocr_qwen_vision.py 201-205")
        sys.exit(1)
    
    batch = sys.argv[1]
    base_dir = "/opt/openclaw/config/workspace-main/Chuan Muc Sadi"
    pdf_path = f"{base_dir}/pdf/Sadi-{batch}.pdf"
    output_path = f"{base_dir}/extracted/Sadi-{batch}-qwen.md"
    temp_dir = f"{base_dir}/extracted/temp_{batch}"
    
    print(f"🔍 Bắt đầu OCR batch Sadi-{batch}")
    
    if not os.path.exists(pdf_path):
        print(f"❌ File PDF không tồn tại: {pdf_path}")
        sys.exit(1)
    
    # Chuyển PDF sang images
    print(f"📄 Chuyển PDF sang hình ảnh...")
    Path(temp_dir).mkdir(parents=True, exist_ok=True)
    
    try:
        images = convert_from_path(pdf_path, dpi=300)
        print(f"  Đã chuyển {len(images)} trang")
        
        image_paths = []
        for i, image in enumerate(images):
            img_path = os.path.join(temp_dir, f"page_{i+1:03d}.png")
            image.save(img_path, 'PNG')
            image_paths.append(img_path)
    except Exception as e:
        print(f"❌ Lỗi chuyển PDF: {e}")
        sys.exit(1)
    
    # OCR với Qwen Vision
    print(f"🤖 OCR với Qwen Vision...")
    success = ocr_with_qwen_vision(image_paths, output_path)
    
    # Dọn dẹp
    import shutil
    try:
        shutil.rmtree(temp_dir)
        print("  Đã dọn file tạm")
    except:
        pass
    
    if success:
        print(f"✅ Hoàn thành! Kết quả: {output_path}")
    else:
        print("❌ OCR thất bại")
        sys.exit(1)

if __name__ == "__main__":
    main()