#!/usr/bin/env python3
"""
OCR Sadi-176-180 với Gemini 2.5 Pro - Simple version
"""

import base64
import os
import urllib.request
import urllib.error
import json
import time

API_KEY = os.environ.get("GEMINI_API_KEY")
WORKDIR = "/opt/openclaw/.openclaw/workspace/Chuan Muc Sadi"
GEMINI_URL = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro:generateContent?key={API_KEY}"

print("📖 Bắt đầu OCR với Gemini 2.5 Pro...")
start_time = time.time()

output_lines = ["# Sadi-176-180 (Raw OCR - Gemini 2.5 Pro)", ""]

for i in range(1, 6):
    print(f"\n🔄 Xử lý trang {i}/5...")
    
    img_path = f"{WORKDIR}/sadi_page-{i}.png"
    
    # Đọc và encode ảnh
    with open(img_path, "rb") as f:
        img_data = base64.b64encode(f.read()).decode("utf-8")
    print(f"   Đã encode trang {i} ({len(img_data)} chars)")
    
    # Tạo payload
    payload = {
        "contents": [{
            "parts": [
                {"text": "Extract all Myanmar text from this page EXACTLY as written. This is a literal transcription - do NOT clean or modify anything. Output ONLY Myanmar text."},
                {"inline_data": {"mime_type": "image/png", "data": img_data}}
            ]
        }],
        "generationConfig": {
            "temperature": 0.1,
            "maxOutputTokens": 4096
        }
    }
    
    # Gửi request
    print("   Đang gửi request...")
    req = urllib.request.Request(
        GEMINI_URL,
        data=json.dumps(payload).encode('utf-8'),
        headers={'Content-Type': 'application/json'},
        method='POST'
    )
    
    try:
        with urllib.request.urlopen(req, timeout=180) as response:
            result = json.loads(response.read().decode('utf-8'))
            text = result["candidates"][0]["content"]["parts"][0]["text"]
            output_lines.append(f"--- Trang {i} ---")
            output_lines.append(text)
            output_lines.append("")
            print(f"   ✅ Trang {i} hoàn thành: {len(text)} ký tự")
    except Exception as e:
        print(f"   ❌ Lỗi trang {i}: {e}")
        output_lines.append(f"--- Trang {i} ---")
        output_lines.append(f"[OCR FAILED: {e}]")
        output_lines.append("")
    
    time.sleep(2)

# Lưu kết quả
output_path = f"{WORKDIR}/extracted/Sadi-176-180-Gemini.md"
with open(output_path, "w", encoding="utf-8") as f:
    f.write("\n".join(output_lines))

elapsed = time.time() - start_time
print(f"\n✅ Hoàn thành! Kết quả lưu tại: {output_path}")
print(f"⏱️ Thời gian: {elapsed:.1f} giây")