#!/usr/bin/env python3
"""Translate sample-source.ass three ways with a local OpenAI-compatible model.

Conditions:
  whole-file   paste the entire .ass file and ask for the translated file back
  text-only    send only the Text field of each Dialogue line as JSON, tags left inline
  protected    same as text-only, but every {...} block, \\N, \\n and \\h is replaced
               with a numbered placeholder first and restored afterwards

Usage:
    python3 run_llm_ass_test.py --runs 3
Outputs: out-<condition>-<run>.ass next to this script, plus raw model replies.

Written by BookTranslator (https://www.booktranslator.app). Free to reuse.
"""

import argparse
import json
import re
import urllib.request
from pathlib import Path

HERE = Path(__file__).parent
URL = "http://localhost:1234/v1/chat/completions"
MODEL = "qwen/qwen3.6-27b"
TARGET = "Simplified Chinese"
TOKEN = re.compile(r"\{[^}]*\}|\\N|\\n|\\h")


def ask(prompt, temperature):
    body = {
        "model": MODEL,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": temperature,
        "reasoning_effort": "none",
    }
    req = urllib.request.Request(URL, json.dumps(body).encode(), {"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=900) as r:
        return json.load(r)["choices"][0]["message"]["content"]


def strip_fence(text):
    m = re.search(r"```(?:\w+)?\n(.*?)```", text, re.S)
    return m.group(1) if m else text


def dialogue_lines(lines):
    """Index and Text field of every Dialogue line that is not a vector drawing."""
    out = []
    for i, line in enumerate(lines):
        if line.startswith("Dialogue:"):
            text = line.split(",", 9)[9]
            if not re.search(r"\\p[1-9]", text):
                out.append((i, text))
    return out


def whole_file(source, temperature):
    prompt = (f"Translate the following ASS subtitle file from English into {TARGET}. "
              "Return the complete file and keep the format.\n\n" + source)
    reply = ask(prompt, temperature)
    return strip_fence(reply), reply


def text_only(source, temperature, protect):
    lines = source.splitlines()
    items, saved = [], {}
    for idx, text in dialogue_lines(lines):
        if protect:
            tokens = TOKEN.findall(text)
            saved[idx] = tokens
            n = iter(range(1, len(tokens) + 1))
            text = TOKEN.sub(lambda m: f"⟦{next(n)}⟧", text)
        items.append({"id": idx, "text": text})
    rule = ("Keep every ⟦number⟧ placeholder exactly once, in a sensible position."
            if protect else
            "Keep every {...} override block and every \\N, \\n and \\h exactly as written.")
    prompt = (f"Translate the \"text\" of each item from English into {TARGET}. {rule} "
              "Return only a JSON array of objects with the same \"id\" values and translated \"text\".\n\n"
              + json.dumps(items, ensure_ascii=False, indent=1))
    reply = ask(prompt, temperature)
    translated = {o["id"]: o["text"] for o in json.loads(strip_fence(reply))}
    for idx, _ in dialogue_lines(lines):
        new = translated.get(idx, "")
        if protect:
            tokens = saved[idx]
            new = re.sub(r"⟦(\d+)⟧", lambda m: tokens[int(m.group(1)) - 1]
                         if 0 < int(m.group(1)) <= len(tokens) else "", new)
        head = lines[idx].split(",", 9)[:9]
        lines[idx] = ",".join(head + [new])
    return "\n".join(lines) + "\n", reply


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--runs", type=int, default=3)
    ap.add_argument("--temperature", type=float, default=0.3)
    args = ap.parse_args()
    source = (HERE / "sample-source.ass").read_text(encoding="utf-8")
    for run in range(1, args.runs + 1):
        for name in ("whole-file", "text-only", "protected"):
            if name == "whole-file":
                out, raw = whole_file(source, args.temperature)
            else:
                out, raw = text_only(source, args.temperature, protect=(name == "protected"))
            (HERE / f"out-{name}-{run}.ass").write_text(out, encoding="utf-8")
            (HERE / f"raw-{name}-{run}.txt").write_text(raw, encoding="utf-8")
            print("wrote", f"out-{name}-{run}.ass")


if __name__ == "__main__":
    main()