import os
import base64
import time
import requests
import sys
import json

if len(sys.argv) < 3:
    print("Usage: python3 ocr_worker_v2.py <start> <end>")
    sys.exit(1)

start = int(sys.argv[1])
end = int(sys.argv[2])

# Google Gemini API key (đã cấu hình trong OpenClaw)
API_KEY = "AIzaSyAkOSPlxUfhUUpDteCX4mrEJnH06rLALFU"

img_dir = "/home/tuan-nguyen/.openclaw/workspace/vi dieu phap giang giai/img/"
out_dir = "/home/tuan-nguyen/.openclaw/workspace/vi dieu phap giang giai/extracted/"
os.makedirs(out_dir, exist_ok=True)

start_str = f"{start:03d}"
end_str = f"{end:03d}"
out_file = os.path.join(out_dir, f"vdp-{start_str}-{end_str}.md")

files = [f"{i:03d}.jpg" for i in range(start, end + 1)]

# Gemini 3.1 Pro Preview API endpoint
url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-3.1-pro-preview:generateContent?key={API_KEY}"

system_instruction = """Trích xuất toàn bộ văn bản tiếng Myanmar từ file này. 
Bỏ qua header và footer, không cần lấy tiêu đề, số trang trong header và footer
Dùng định dạng markdown để giữ chính xác cách trình bày của file gốc: tô đậm, in nghiêng, lùi dòng v.v…
Giữa các đoạn nên có 1 dòng trống cho dễ đọc
Đây là văn bản Phật pháp nên cần chính xác tuyệt đối, đặc biệt với thuật ngữ Pali và số Myanmar"""

with open(out_file, "w", encoding="utf-8") as f:
    for filename in files:
        filepath = os.path.join(img_dir, filename)
        if not os.path.exists(filepath):
            print(f"File not found: {filepath}")
            continue
            
        with open(filepath, "rb") as image_file:
            base64_image = base64.b64encode(image_file.read()).decode('utf-8')
        
        page_num = filename.split('.')[0]
        
        payload = {
            "system_instruction": {
                "parts": [{"text": system_instruction}]
            },
            "contents": [
                {
                    "role": "user",
                    "parts": [
                        {"text": f"Trang pdf {page_num}"},
                        {
                            "inline_data": {
                                "mime_type": "image/jpeg",
                                "data": base64_image
                            }
                        }
                    ]
                }
            ],
            "generationConfig": {
                "temperature": 0.1,
                "maxOutputTokens": 65536
            }
        }
        
        print(f"Processing {filename}...")
        
        response = requests.post(url, headers={"Content-Type": "application/json"}, json=payload)
        
        if response.status_code == 200:
            data = response.json()
            try:
                content = data['candidates'][0]['content']['parts'][0]['text']
                f.write(content + "\n\n---\n\n")
                print(f"Success for {filename}")
            except (KeyError, IndexError) as e:
                print(f"Failed to parse response for {filename}: {data}")
                f.write(f"[[LỖI: Không parse được response]]\n\n---\n\n")
        else:
            print(f"Failed for {filename}: {response.status_code} - {response.text[:500]}")
        
        time.sleep(2)

print(f"Done! Saved to {out_file}")