#!/usr/bin/env python3
"""Compare a translated LaTeX file with its source and report structural damage.

Usage:
    python3 check_latex_translation.py source.tex translated.tex [--keep TERM ...]

It checks the parts of a .tex file that must survive translation unchanged:
the preamble, command names, environments, labels, references, citations,
math, verbatim blocks, URLs, ties before references, and LaTeX special
characters. It does not judge translation quality and does not compile the
file; compile both files and compare the PDFs as well.

Exit code 0 when no FAIL lines are printed, 1 otherwise.

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

import argparse
import re
import sys
from collections import Counter

VERBATIM_ENVS = ("verbatim", "lstlisting", "minted")
MATH_ENVS = ("equation", "equation*", "align", "align*", "gather", "gather*",
             "multline", "multline*", "displaymath", "math")
REF_CMDS = ("ref", "eqref", "pageref", "autoref", "cref", "Cref", "nameref")
CITE_CMDS = ("cite", "citep", "citet", "parencite", "textcite", "autocite")
# Commands whose braced argument is reader-facing text inside math.
MATH_TEXT = re.compile(r"\\(?:text|textrm|textit|textbf|mbox|intertext)\{[^{}]*\}")


def strip_comments(tex):
    """Remove % comments (not \\%) so prose checks ignore them."""
    return re.sub(r"(?<!\\)%.*", "", tex)


def comments(tex):
    return [m.group(1).strip() for m in re.finditer(r"(?<!\\)%(.*)", tex)]


def split_verbatim(tex):
    """Return (tex with verbatim bodies removed, list of verbatim bodies)."""
    bodies = []
    pattern = re.compile(
        r"\\begin\{(" + "|".join(VERBATIM_ENVS) + r")\}(.*?)\\end\{\1\}", re.S)

    def keep(m):
        bodies.append(m.group(2))
        return "\\begin{%s}\\end{%s}" % (m.group(1), m.group(1))

    return pattern.sub(keep, tex), bodies


def preamble(tex):
    head = tex.split("\\begin{document}", 1)[0]
    # \title, \author and \date hold reader-facing text; mask their arguments.
    return re.sub(r"\\(title|author|date)\{.*?\}\s*$", r"\\\1{...}", head,
                  flags=re.S | re.M)


def math_segments(tex):
    envs = "|".join(re.escape(e) for e in MATH_ENVS)
    pattern = re.compile(
        r"\\begin\{(" + envs + r")\}(.*?)\\end\{\1\}"
        r"|\\\[(.*?)\\\]"
        r"|\\\((.*?)\\\)"
        r"|(?<!\\)\$(.+?)(?<!\\)\$", re.S)
    out = []
    for m in pattern.finditer(tex):
        body = next(g for g in m.groups()[1:] if g is not None)
        body = MATH_TEXT.sub(lambda t: t.group(0).split("{")[0] + "{...}", body)
        out.append(re.sub(r"\s+", " ", body).strip())
    return out


def keys(tex, cmds):
    found = []
    for m in re.finditer(r"\\(" + "|".join(cmds) + r")\*?(?:\[[^\]]*\])*\{([^}]*)\}", tex):
        found += [k.strip() for k in m.group(2).split(",")]
    return Counter(found)


def command_names(tex):
    return Counter(re.findall(r"\\([A-Za-z]+)", tex))


def environments(tex):
    return re.findall(r"\\(begin|end)\{([^}]*)\}", tex)


def urls(tex):
    return Counter(re.findall(r"\\(?:url|href)\{([^}]*)\}", tex))


def ties(tex):
    return len(re.findall(r"~\\(?:" + "|".join(REF_CMDS + CITE_CMDS) + r")\b", tex))


def unescaped(tex, char):
    return len(re.findall(r"(?<!\\)" + re.escape(char), tex))


def brace_balance(tex):
    depth = 0
    for ch in re.sub(r"\\[{}]", "", tex):
        depth += ch == "{"
        depth -= ch == "}"
        if depth < 0:
            return -1
    return depth


def diff_counter(a, b):
    missing = a - b
    extra = b - a
    parts = []
    if missing:
        parts.append("missing " + ", ".join(f"{k} x{v}" for k, v in sorted(missing.items())))
    if extra:
        parts.append("added " + ", ".join(f"{k} x{v}" for k, v in sorted(extra.items())))
    return "; ".join(parts)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("source")
    ap.add_argument("translated")
    ap.add_argument("--keep", nargs="*", default=[],
                    help="terms that must appear as often in the translation, e.g. product names")
    args = ap.parse_args()
    src_raw = open(args.source, encoding="utf-8").read()
    dst_raw = open(args.translated, encoding="utf-8").read()

    src_nv, src_verb = split_verbatim(src_raw)
    dst_nv, dst_verb = split_verbatim(dst_raw)
    src = strip_comments(src_nv)
    dst = strip_comments(dst_nv)
    failures = 0

    def report(level, name, detail=""):
        nonlocal failures
        failures += level == "FAIL"
        print(f"{level:4} {name}" + (f": {detail}" if detail else ""))

    report("PASS" if preamble(src) == preamble(dst) else "FAIL", "preamble unchanged",
           "" if preamble(src) == preamble(dst) else "packages, macros or options differ")

    d = diff_counter(command_names(src), command_names(dst))
    report("FAIL" if d else "PASS", "command names", d)

    se, de = environments(src), environments(dst)
    report("PASS" if se == de else "FAIL", "environments in the same order",
           "" if se == de else diff_counter(Counter(se), Counter(de)) or "order changed")

    for name, cmds in (("labels", ("label",)), ("references", REF_CMDS),
                       ("citations", CITE_CMDS), ("bibliography keys", ("bibitem",))):
        d = diff_counter(keys(src, cmds), keys(dst, cmds))
        report("FAIL" if d else "PASS", name, d)

    # Compared as a multiset: translation may legitimately reorder inline math.
    sm, dm = math_segments(src), math_segments(dst)
    d = diff_counter(Counter(sm), Counter(dm))
    report("FAIL" if d else "PASS", "math unchanged",
           d or f"{len(sm)} segments; \\text{{}} content may be translated")

    report("PASS" if src_verb == dst_verb else "FAIL", "verbatim blocks unchanged",
           f"{len(src_verb)} blocks")

    d = diff_counter(urls(src), urls(dst))
    report("FAIL" if d else "PASS", "URLs", d)

    for char, label in (("%", "unescaped % (starts a comment)"),
                        ("&", "unescaped & (column separator)"),
                        ("#", "unescaped # (macro parameter)")):
        a = unescaped(src_nv, char) if char != "%" else len(comments(src_nv))
        b = unescaped(dst_nv, char) if char != "%" else len(comments(dst_nv))
        report("PASS" if a == b else "FAIL", label, f"{a} in source, {b} in translation")

    a, b = brace_balance(src), brace_balance(dst)
    report("PASS" if a == b == 0 else "FAIL", "braces balanced",
           "" if a == b == 0 else f"depth {b} at end of translation")

    a, b = ties(src), ties(dst)
    report("PASS" if a == b else "WARN", "ties (~) before \\ref and \\cite",
           f"{a} in source, {b} in translation")

    sc, dc = comments(src_nv), comments(dst_nv)
    if sc == dc:
        report("PASS", "comments left as written", f"{len(sc)} comments")
    elif len(sc) == len(dc):
        report("WARN", "comments translated", "decide whether readers of the source need them")
    else:
        report("WARN", "comments added or removed", f"{len(sc)} in source, {len(dc)} in translation")
    for term in args.keep:
        a, b = src_raw.count(term), dst_raw.count(term)
        report("PASS" if a == b else "FAIL", f"keep '{term}'", f"{a} in source, {b} in translation")

    sys.exit(1 if failures else 0)


if __name__ == "__main__":
    main()