#!/usr/bin/env python3
"""Translate sample-source.mdx 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
  protected    replace front matter keys, code, URLs, JSX, comments, footnote
               labels and admonition markers with numbered placeholders, translate
               the rest, then restore the originals

Usage:
    python3 run_llm_markdown_test.py --runs 3
Outputs: out-<condition>-<run>.mdx next to this script; for the protected condition it
prints any placeholder the model did not return.

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:
- In the YAML front matter, keep every key unchanged. Translate only the values of title and description.
- Do not change anything inside fenced code blocks or `inline code`.
- Do not change URLs, link targets, image paths, or reference definitions such as [glossary]: ./glossary.md.
- Keep import lines, JSX tags and their attribute names, and {/* comments */} unchanged.
- Keep footnote labels such as [^sync], admonition markers such as :::tip, and table structure.
- Translate headings, paragraphs, list items, table cells, link text, image alt text, and the title attribute of JSX tags."""

# Order matters: larger spans first so inner patterns are not masked twice.
PROTECT = re.compile(
    r"^```[\s\S]*?^```"                      # fenced code blocks
    r"|^import .*$"                          # MDX imports
    r"|\{/\*[\s\S]*?\*/\}"                   # MDX comments
    r"|^:::\w*$"                             # admonition markers
    r"|^\[[^\]^]+\]:\s*\S+.*$"               # reference definitions
    r"|\[\^[^\]]+\]:?"                       # footnote labels
    r"|`[^`\n]+`"                            # inline code
    r"|\]\([^)\s]+(?:\s+\"[^\"]*\")?\)"      # link and image targets
    r"|\]\[[^\]]+\]"                         # reference link labels
    r"|</?[A-Za-z][^>]*>",                   # JSX and HTML tags
    re.M,
)
FRONT = re.compile(r"\A---\n(.*?)\n---\n", re.S)


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


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


def protected(source, temperature):
    saved = []

    def keep(m):
        saved.append(m.group(0))
        return f"⟦{len(saved) - 1}⟧"

    # Front matter: protect everything except the title and description values.
    fm = FRONT.match(source)
    body = source[fm.end():]
    fm_lines = []
    for line in fm.group(1).split("\n"):
        k = re.match(r'^(title|description): "(.*)"$', line)
        if k:
            saved.append(f'{k.group(1)}: "')
            fm_lines.append(f"⟦{len(saved) - 1}⟧{k.group(2)}\"")
        else:
            saved.append(line)
            fm_lines.append(f"⟦{len(saved) - 1}⟧")
    masked = "---\n" + "\n".join(fm_lines) + "\n---\n" + PROTECT.sub(keep, body)

    prompt = (f"Translate the following Markdown text from English into {TARGET}. "
              "Tokens like ⟦12⟧ are placeholders: keep every placeholder exactly once, "
              "in a natural position, and do not translate or change them. "
              "Return only the translated text.\n\n" + masked)
    reply = strip_fence(ask(prompt, temperature))
    missing = [i for i in range(len(saved)) if f"⟦{i}⟧" not in reply]
    out = re.sub(r"⟦(\d+)⟧", lambda m: saved[int(m.group(1))], reply)
    return out, missing


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--runs", type=int, default=3)
    ap.add_argument("--start", type=int, default=1)
    ap.add_argument("--temperature", type=float, default=0.3)
    args = ap.parse_args()
    source = (HERE / "sample-source.mdx").read_text(encoding="utf-8")
    for run in range(args.start, args.start + args.runs):
        for name in ("whole-file", "instructed", "protected"):
            if name == "protected":
                out, missing = protected(source, args.temperature)
                note = f"missing placeholders: {missing}" if missing else "all placeholders returned"
            else:
                out = whole_file(source, args.temperature, rules=(name == "instructed"))
                note = ""
            path = HERE / f"out-{name}-{run}.mdx"
            path.write_text(out, encoding="utf-8")
            print(f"{path.name}\t{note}", flush=True)


if __name__ == "__main__":
    main()