import os, time, re, polib, openai
from tqdm import tqdm

SRC_PO = "iron-security.pot"
DST_PO = "iron-security-lt_LT.po"
MODEL  = "gpt-4o-mini"
BATCH  = 30
SLEEP  = 1
DELIM  = " ||| "            # saugus skyriklis

openai.api_key = os.getenv("OPENAI_API_KEY")

def protect(text):
    return re.sub(r"(%\d*\$?[sd]|%\d*\$?f|%\d*\$?\.?\d*[ds])", r"⟦\1⟧", text)

po   = polib.pofile(SRC_PO)
todo = [e for e in po if not e.translated() and e.msgid.strip()]

for i in tqdm(range(0, len(todo), BATCH), desc="Translating"):
    chunk = todo[i:i+BATCH]

    # ---- paruošiame žinutę ----
    user_lines = [
        f"{idx}:{protect(e.msgid)}"
        for idx, e in enumerate(chunk, start=1)
    ]
    prompt = (
        "Išversk kiekvieną eilutę į lietuvių kalbą. "
        f"Grąžink TIK vertimus, ta pačia tvarka, atskirtus '{DELIM}'. "
        "NEpridėk jokių papildomų paaiškinimų."
        "\n\n" + "\n".join(user_lines)
    )

    resp = openai.chat.completions.create(
        model=MODEL,
        messages=[{"role": "system",
                   "content": "You are a professional Lithuanian software translator;"\
                              " keep placeholders (⟦%s⟧) identiški."},
                  {"role": "user", "content": prompt}],
        temperature=0.2,
    )

    translations = resp.choices[0].message.content.split(DELIM)
    translations = [re.sub(r"⟦(.*?)⟧", r"\1", t.strip()) for t in translations]

    if len(translations) != len(chunk):
        raise ValueError(
            f"Expected {len(chunk)} vertimų, gauta {len(translations)} – "
            "patikrinkite atsakymą arba sumažinkite BATCH."
        )

    for entry, lt in zip(chunk, translations):
        entry.msgstr = lt or entry.msgid

    time.sleep(SLEEP)

po.save(DST_PO)
po.save_as_mofile(DST_PO.replace(".po", ".mo"))
print("✓ Vertimai užbaigti →", DST_PO)