import os
import sys
import time
import subprocess
import json
import urllib.request

api_key = os.environ.get("GEMINI_API_KEY")

ranges = [(311, 320), (321, 330), (331, 340), (341, 350)]

prompt_template = """Bạn là một trợ lý dịch thuật xuất sắc chuyên về Phật giáo Nguyên thủy (Theravada).
Nhiệm vụ của bạn là dịch văn bản sau sang tiếng Việt (sử dụng đúng thuật ngữ Theravada).
Giữ nguyên cấu trúc Markdown. Ở đầu mỗi trang dịch, chú thích "Trang [số trang]".
Giữa các đoạn văn cách nhau một dòng trống.

Nội dung:
{text}
"""

def generate_content(prompt):
    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent?key={api_key}"
    headers = {"Content-Type": "application/json"}
    data = {"contents": [{"parts": [{"text": prompt}]}]}
    req = urllib.request.Request(url, data=json.dumps(data).encode("utf-8"), headers=headers)
    try:
        with urllib.request.urlopen(req) as response:
            res_json = json.loads(response.read().decode("utf-8"))
            return res_json["candidates"][0]["content"]["parts"][0]["text"]
    except Exception as e:
        print(f"Error: {e}")
        return "Error"

for start, end in ranges:
    print(f"Translating {start}-{end}...")
    subprocess.run(["pdftotext", "-f", str(start), "-l", str(end), "buddhistmonasticcode-1.pdf", "tmp.txt"])
    with open("tmp.txt", "r") as f:
        text = f.read()
    
    pages = text.split('\x0c')
    formatted_text = ""
    for i, p_text in enumerate(pages):
        if p_text.strip() and (start + i) <= end:
            formatted_text += f"\nTrang {start + i}\n\n{p_text.strip()}\n\n"

    prompt = prompt_template.format(text=formatted_text)
    
    translated_text = generate_content(prompt)
    
    with open(f"buddhistmonasticcode_vi_{start}_{end}.md", "w") as f:
        f.write(translated_text)
    
    print(f"Done {start}-{end}. Sleeping for 1200s...")
    sys.stdout.flush()
    time.sleep(1200)

print("All translations completed.")