#!/usr/bin/env python3
"""
OCR Sadi-176-180 với Gemini 2.5 Pro - Gửi từng trang
"""

import base64
import os
import requests
from PIL import Image
import io
import time

start_time = time.time()

API_KEY = os.environ.get("GEMINI_API_KEY")
if not API_KEY:
    print("❌ Thiếu GEMINI_API_KEY")
    exit(1)

WORKDIR = "/opt/openclaw/.openclaw/workspace/Chuan Muc Sadi"

all_text = []

# Gemini 2.5 Pro endpoint
GEMINI_URL = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro:generateContent?key={API_KEY}"

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 Image.open(img_path) as img:
        # Resize để giảm kích thước nhưng vẫn giữ chất lượng
        ratio = 1600 / img.width
        new_height = int(img.height * ratio)
        img_resized = img.resize((1600, 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}")
    
    # Prompt cho Myanmar OCR
    prompt = """Extract all Myanmar text from this page EXACTLY as written. This is a BUDDHIST TEXT transcription task.

CRITICAL RULES:
1. Preserve EVERY character literally - do NOT clean, modify, or autocorrect ANYTHING
2. Keep all diacritics, vowels, consonants exactly as they appear
3. Keep repeated characters if they exist in the image
4. Keep any apparent errors - this is literal transcription
5. Do NOT add English explanations or translations
6. Pay special attention to:
   - Numbers: ၀ ၁ ၂ ၃ ၄ ၅ ၆ ၇ ၈ ၉ (especially ၈ vs ဂ)
   - Proper names: Channa, Sakka, အာဠာဝက
   - Particles: လျှင်, သူတစ်ပါး
   - Compound words: ဦးခေါင်း, မျက်ခုံး, လက်ယာရစ်

Output ONLY the Myanmar text with page markers."""

    # Gemini API payload
    payload = {
        "contents": [{
            "parts": [
                {"text": prompt},
                {"inline_data": {"mime_type": "image/png", "data": img_data}}
            ]
        }],
        "generationConfig": {
            "temperature": 0.1,
            "topP": 0.8,
            "topK": 40,
            "maxOutputTokens": 4096
        }
    }
    
    headers = {
        "Content-Type": "application/json"
    }
    
    try:
        response = requests.post(
            GEMINI_URL,
            headers=headers,
            json=payload,
            timeout=180
        )
        
        if response.status_code == 200:
            result = response.json()
            if "candidates" in result and len(result["candidates"]) > 0:
                text = result["candidates"][0]["content"]["parts"][0]["text"]
                all_text.append(f"--- Trang {i} ---\n{text}")
                print(f"   ✅ Trang {i} hoàn thành: {len(text)} ký tự")
            else:
                print(f"   ❌ Không có kết quả trang {i}")
                all_text.append(f"--- Trang {i} ---\n[NO RESPONSE]")
        else:
            print(f"   ❌ Lỗi trang {i}: {response.status_code}")
            print(f"   {response.text[:300]}")
            all_text.append(f"--- Trang {i} ---\n[OCR FAILED: {response.status_code}]")
            
    except Exception as e:
        print(f"   ❌ Exception trang {i}: {e}")
        all_text.append(f"--- Trang {i} ---\n[OCR EXCEPTION: {str(e)}]")
    
    # Nghỉ một chút giữa các trang
    time.sleep(1)

# Lưu kết quả thô
if all_text:
    output_path = f"{WORKDIR}/extracted/Sadi-176-180-Gemini.md"
    with open(output_path, "w", encoding="utf-8") as f:
        f.write("# Sadi-176-180 (Raw OCR - Gemini 2.5 Pro)\n\n")
        f.write("\n\n".join(all_text))
    
    elapsed = time.time() - start_time
    print(f"\n✅ OCR hoàn thành! Đã lưu vào: {output_path}")
    total_chars = sum(len(t) for t in all_text)
    print(f"📊 Tổng: {total_chars} ký tự")
    print(f"⏱️ Thời gian: {elapsed:.1f} giây")
else:
    print("\n❌ Không có kết quả OCR nào!")