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

Conditions:
  whole-file   paste the entire file and ask for the translated file back
  instructed   the same, plus explicit rules about what must not change
  segments     keep the preamble out of the request (only the \\title line is
               sent), send each blank-line-separated block of the body on its own
               with the same rules, skip verbatim blocks, and reassemble

Usage:
    python3 run_llm_latex_test.py --runs 3
Outputs: out-<condition>-<run>.tex next to this script.

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"

RULES = """Rules:
- Do not change the preamble: \\documentclass, \\usepackage lines and \\newcommand definitions. Translate only the text inside \\title{}.
- Keep every command, environment name, \\label, \\ref, \\eqref, \\cite and \\bibitem key exactly as written.
- Do not change anything inside math ($...$, \\[...\\], equation), except text inside \\text{}.
- Do not change verbatim blocks, \\url{} arguments or the first argument of \\href.
- Keep LaTeX special characters escaped: write \\% \\& \\$ \\# in text, never bare % & $ #.
- Keep ~ before \\ref and \\cite.
- Translate the abstract, paragraphs, headings, list items, captions, footnotes, table text and the second argument of \\href."""


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):
    """Remove one outer ``` fence if the model wrapped the whole reply in it."""
    m = re.match(r"\s*```(?:latex|tex)?\n(.*?)\n?```\s*\Z", text, re.S)
    return m.group(1) if m else text.strip("\n")


def whole_file(source, temperature, rules=False):
    prompt = (f"Translate the following LaTeX file from English into {TARGET}. "
              "Return the complete file and keep it compilable.\n\n")
    if rules:
        prompt += RULES + "\n\n"
    return strip_fence(ask(prompt + source, temperature)) + "\n"


def segments(source, temperature):
    head, rest = source.split("\\begin{document}", 1)
    body, tail = rest.rsplit("\\end{document}", 1)
    prompt = (f"Translate this fragment of a LaTeX file from English into {TARGET}. "
              "Return only the translated fragment, with the same line structure.\n\n"
              + RULES + "\n\nFragment:\n")
    title = re.search(r"^\\title\{.*\}$", head, re.M)
    head = head.replace(title.group(0), strip_fence(ask(prompt + title.group(0), temperature)))
    # Blank lines separate blocks; verbatim environments are split out and never sent.
    blocks = re.split(r"(\n\s*\n|\\begin\{verbatim\}.*?\\end\{verbatim\}\n)", body, flags=re.S)
    out = []
    for block in blocks:
        if not re.search(r"[A-Za-z]{3}", re.sub(r"\\[A-Za-z]+|\{[^{}]*\}", "", block)) \
                or block.startswith("\\begin{verbatim}"):
            out.append(block)
            continue
        out.append(strip_fence(ask(prompt + block.strip("\n"), temperature)))
        # keep the original leading/trailing newlines of the block
        lead = re.match(r"\n*", block).group(0)
        trail = re.search(r"\n*\Z", block).group(0)
        out[-1] = lead + out[-1].strip("\n") + trail
    return head + "\\begin{document}" + "".join(out) + "\\end{document}" + tail


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.tex").read_text(encoding="utf-8")
    for run in range(1, args.runs + 1):
        for name, fn in (("whole-file", lambda s: whole_file(s, args.temperature)),
                         ("instructed", lambda s: whole_file(s, args.temperature, rules=True)),
                         ("segments", lambda s: segments(s, args.temperature))):
            path = HERE / f"out-{name}-{run}.tex"
            path.write_text(fn(source), encoding="utf-8")
            print("wrote", path.name)


if __name__ == "__main__":
    main()