#!/usr/bin/env python3
"""Translate sample-source.html 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     extract text nodes and translatable attributes with a tokenizer,
               send them as a numbered JSON list, then put the translations back
               into the untouched markup

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

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

import argparse
import html
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"
LANG = "zh-Hans"

RULES = """Rules:
- Keep every tag, attribute name, and the order of elements unchanged.
- Do not change href, src, id, class, for, name, type, action, method, width, height, or data-* values.
- Translate the text people read: the title element, headings, paragraphs, table cells, labels, link text,
  and the values of alt, title, placeholder, and aria-label, plus the content of the description and og:title meta tags.
- Do not change anything inside <code>, <pre>, <script>, or <style>, except the name and description values in the JSON-LD block.
- Do not translate text inside an element with translate="no".
- Keep HTML entities such as &lt; &gt; &amp; &nbsp; as they are.
- Keep HTML comments unchanged.
- Change the lang attribute of the html element to zh-Hans."""

TOKEN = re.compile(r"<!--[\s\S]*?-->|<![^>]*>|<[^>]+>|[^<]+")
TAG_NAME = re.compile(r"^</?\s*([a-zA-Z0-9]+)")
SKIP = {"script", "style", "code", "pre"}
ATTRS = re.compile(r'\b(alt|title|placeholder|aria-label)="([^"]*)"')


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.match(r"\s*```(?:html)?\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 HTML file from English into {TARGET}. "
              "Return the complete file and keep the HTML markup.\n\n")
    if rules:
        prompt += RULES + "\n\n"
    return strip_fence(ask(prompt + source, temperature))


def segments(source, temperature):
    tokens = TOKEN.findall(source)
    texts, slots = [], []  # slots: (token index, kind, extra)
    stack = []  # open elements whose text we skip
    for i, tok in enumerate(tokens):
        if tok.startswith("<!--") or tok.startswith("<!"):
            continue
        if tok.startswith("<"):
            m = TAG_NAME.match(tok)
            name = m.group(1).lower() if m else ""
            closing = tok.startswith("</")
            if closing:
                if stack and stack[-1] == name:
                    stack.pop()
                continue
            if name in SKIP or 'translate="no"' in tok:
                stack.append(name)
            for am in ATTRS.finditer(tok):
                slots.append((i, "attr", am.group(1)))
                texts.append(html.unescape(am.group(2)))
            if name == "meta" and re.search(r'(name="description"|property="og:title")', tok):
                slots.append((i, "meta", None))
                texts.append(html.unescape(re.search(r'content="([^"]*)"', tok).group(1)))
            continue
        if stack or not tok.strip():
            continue
        slots.append((i, "text", None))
        texts.append(html.unescape(tok))

    prompt = (f"Translate each string in this JSON array from English into {TARGET}. "
              "Return only a JSON array of the same length, in the same order.\n\n"
              + json.dumps(texts, ensure_ascii=False))
    reply = strip_fence(ask(prompt, temperature)).strip()
    reply = re.sub(r"^```(?:json)?\s*|\s*```$", "", reply)
    out = json.loads(reply)
    if len(out) != len(texts):
        raise SystemExit(f"model returned {len(out)} strings for {len(texts)}")

    for (i, kind, extra), new in zip(slots, out):
        tok = tokens[i]
        if kind == "text":
            lead = tok[: len(tok) - len(tok.lstrip())]
            trail = tok[len(tok.rstrip()):]
            tokens[i] = lead + html.escape(new.strip(), quote=False) + trail
        elif kind == "attr":
            tokens[i] = re.sub(rf'\b{extra}="[^"]*"', lambda m: f'{extra}="{html.escape(new)}"', tok, count=1)
        else:
            tokens[i] = re.sub(r'content="[^"]*"', lambda m: f'content="{html.escape(new)}"', tok, count=1)
    result = "".join(tokens)
    return re.sub(r'<html lang="[^"]*"', f'<html lang="{LANG}"', result, count=1)


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


if __name__ == "__main__":
    main()