#!/usr/bin/env python3
"""
OCR Sadi-176-180 với Gemini 2.5 Pro - Với resize ảnh
"""

import base64
import os
import urllib.request
import json
import time
from PIL import Image
import io

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...", flush=True)
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...", flush=True)
    
    img_path = f"{WORKDIR}/sadi_page-{i}.png"
    
    # Resize ảnh để giảm kích thước
    with Image.open(img_path) as img:
        # Resize xuống 800px width
        ratio = 800 / img.width
        new_height = int(img.height * ratio)
        img_resized = img.resize((800, new_height), Image.LANCZOS)
        
        buffer = io.BytesIO()
        img_resized.save(buffer, format="PNG", optimize=True)
        img_data = base64.b64encode(buffer.getvalue()).decode("utf-8")
    
    print(f"   Đã encode trang {i} ({len(img_data)} chars)", flush=True)
    
    # 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...", flush=True)
    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ự", flush=True)
    except Exception as e:
        print(f"   ❌ Lỗi trang {i}: {e}", flush=True)
        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}", flush=True)
print(f"⏱️ Thời gian: {elapsed:.1f} giây", flush=True)