import os
import requests
import base64
import time
import json

api_key = os.environ.get("OPENROUTER_API_KEY", "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")
headers = {
    "Authorization": f"Bearer {api_key}",
    "Content-Type": "application/json",
    "HTTP-Referer": "https://openclaw.ai",
    "X-Title": "OpenClaw OCR"
}

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)
out_file = os.path.join(out_dir, "vdp-001-005.md")

files = ["001.jpg", "002.jpg", "003.jpg", "004.jpg", "005.jpg"]

prompt_template = """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
[QUAN TRỌNG] 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
Điền thêm chỉ mục trang vào đầu trang "Trang pdf {page_num}" """

with open(out_file, "w", encoding="utf-8") as f:
    for filename in files:
        filepath = os.path.join(img_dir, filename)
        with open(filepath, "rb") as image_file:
            base64_image = base64.b64encode(image_file.read()).decode('utf-8')
        
        page_num = filename.split('.')[0]
        page_prompt = prompt_template.replace('{page_num}', page_num)
        
        payload = {
            "model": "google/gemini-3.1-pro-preview",
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": page_prompt
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{base64_image}"
                            }
                        }
                    ]
                }
            ]
        }
        
        print(f"Processing {filename}...")
        response = requests.post("https://openrouter.ai/api/v1/chat/completions", headers=headers, json=payload)
        
        if response.status_code == 200:
            data = response.json()
            content = data['choices'][0]['message']['content']
            f.write(content + "\n\n---\n\n")
            print(f"Success for {filename}")
        else:
            print(f"Failed for {filename}: {response.text}")
        
        time.sleep(2)

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