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

Usage:
    python3 check_markdown_translation.py source.mdx translated.mdx [--keep TERM ...]

ERROR lines are changes that break rendering, links, or code.
WARN lines are changes that may be intended, such as a translated code comment.
The script does not judge translation quality. Python standard library only.

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

import argparse
import re
import sys
import unicodedata
from collections import Counter

FRONT = re.compile(r"\A---\n(.*?)\n---\n", re.S)
FENCE = re.compile(r"^(```|~~~)([^\n]*)\n(.*?)^\1[ \t]*$", re.S | re.M)
TRANSLATABLE_KEYS = {"title", "description"}


def split(text):
    m = FRONT.match(text)
    return (m.group(1), text[m.end():]) if m else ("", text)


def front_matter(fm):
    pairs = []
    for line in fm.split("\n"):
        m = re.match(r"^([A-Za-z_][\w-]*):\s*(.*)$", line)
        if m:
            pairs.append((m.group(1), m.group(2)))
    return pairs


def fences(body):
    return [(m.group(2).strip(), m.group(3)) for m in FENCE.finditer(body)]


def prose(body):
    """Body with fenced code removed and inline code blanked, for the inline checks."""
    body = FENCE.sub("", body)
    body = re.sub(r"\{/\*.*?\*/\}|<!--.*?-->", "", body, flags=re.S)
    return re.sub(r"`[^`\n]+`", "``", body)


def inline_code(body):
    return Counter(re.findall(r"`[^`\n]+`", FENCE.sub("", body)))


def link_targets(text):
    return Counter(m.group(1) for m in re.finditer(r"\]\(([^)\s]+)", text))


def link_titles(text):
    return re.findall(r'\]\([^)\s]+\s+"([^"]*)"\)', text)


def ref_definitions(text):
    return sorted(re.findall(r"^\[([^\]^]+)\]:\s*(\S+)", text, re.M))


def ref_uses(text):
    return Counter(re.findall(r"\]\[([^\]]+)\]", text))


def headings(text):
    return [(len(m.group(1)), m.group(2).strip()) for m in re.finditer(r"^(#{1,6})\s+(.*)$", text, re.M)]


EXPLICIT_ID = re.compile(r"\s*(?:\{#([^}]+)\}|\{/\*\s*#(\S+?)\s*\*/\}|<!--\s*#(\S+?)\s*-->)\s*$")


def slug(heading):
    """github-slugger style: lowercase, drop punctuation and symbols, spaces to hyphens."""
    heading = EXPLICIT_ID.sub("", heading).lower()
    out = []
    for ch in heading:
        cat = unicodedata.category(ch)
        if ch in "-_" or cat[0] in "LN" or cat == "Mn":
            out.append(ch)
        elif ch == " ":
            out.append("-")
    return "".join(out)


def anchors(text):
    ids = set()
    for _, h in headings(text):
        m = EXPLICIT_ID.search(h)
        ids.add(next(g for g in m.groups() if g) if m else slug(h))
    return ids


def tags(text):
    out = []
    for m in re.finditer(r"<(/?)([A-Za-z][\w.]*)([^>]*)>", text):
        attrs = dict(re.findall(r'([\w-]+)="([^"]*)"', m.group(3)))
        out.append((m.group(1) + m.group(2), attrs))
    return out


def mdx_comments(text):
    return [c for c in re.findall(r"\{/\*(.*?)\*/\}", text, re.S) if not re.fullmatch(r"\s*#\S+\s*", c)]


def tables(text):
    result, rows = [], []
    for line in text.split("\n") + [""]:
        if line.strip().startswith("|"):
            rows.append(line.strip().strip("|").count("|") + 1)
        elif rows:
            result.append(rows)
            rows = []
    return result


def is_punct(ch):
    return ch == "" or unicodedata.category(ch)[0] in "PS"


def is_space(ch):
    return ch == "" or ch.isspace()


def unmatched_emphasis(text):
    """Simplified CommonMark flanking rules: list delimiter runs that cannot open or close."""
    problems = []
    for para in re.split(r"\n\s*\n", text):
        if para.lstrip().startswith(("|", "<", "{")):
            continue
        stack = []
        for m in re.finditer(r"\*+", para):
            if m.start() == 0 or para[m.start() - 1] == "\n":
                rest = para[m.end():m.end() + 1]
                if rest == " ":  # list bullet
                    continue
            before = para[m.start() - 1] if m.start() else ""
            after = para[m.end()] if m.end() < len(para) else ""
            left = not is_space(after) and (not is_punct(after) or is_space(before) or is_punct(before))
            right = not is_space(before) and (not is_punct(before) or is_space(after) or is_punct(after))
            run = m.group(0)
            if right and stack and stack[-1][0] == run:
                stack.pop()
            elif left:
                stack.append((run, m.start(), para))
            else:
                problems.append(context(para, m.start()))
        problems += [context(p, i) for _, i, p in stack]
    return problems


def context(text, i):
    return text[max(0, i - 12):i + 14].replace("\n", " ")


def tildes(text):
    risky = []
    for para in re.split(r"\n\s*\n", text):
        singles = re.findall(r"(?<!~)~(?!~)", para)
        if len(singles) >= 2:
            risky.append(para.strip()[:60].replace("\n", " "))
    return risky


def hard_breaks(text):
    return len(re.findall(r"[^ \n] {2,}\n|\\\n", text))


def check(src_text, out_text, keep):
    errors, warns = [], []
    sfm, sbody = split(src_text)
    ofm, obody = split(out_text)

    s_pairs, o_pairs = front_matter(sfm), front_matter(ofm)
    if [k for k, _ in s_pairs] != [k for k, _ in o_pairs]:
        errors.append(f"front matter keys changed: {[k for k, _ in s_pairs]} -> {[k for k, _ in o_pairs]}")
    o_map = dict(o_pairs)
    for k, v in s_pairs:
        if k in o_map and k not in TRANSLATABLE_KEYS and o_map[k] != v:
            errors.append(f"front matter value '{k}' changed: {v} -> {o_map[k]}")
        if k in TRANSLATABLE_KEYS and o_map.get(k) == v:
            warns.append(f"front matter '{k}' was not translated")

    s_imp = re.findall(r"^(?:import|export) .*$", sbody, re.M)
    o_imp = re.findall(r"^(?:import|export) .*$", obody, re.M)
    if s_imp != o_imp:
        errors.append(f"import/export lines changed: {s_imp} -> {o_imp}")

    sf, of = fences(sbody), fences(obody)
    if [i for i, _ in sf] != [i for i, _ in of]:
        errors.append(f"code blocks changed: {[i for i, _ in sf]} -> {[i for i, _ in of]}")
    else:
        for n, ((info, s), (_, o)) in enumerate(zip(sf, of), 1):
            if s == o:
                continue
            strip = lambda c: re.sub(r"\s*#.*$", "", c, flags=re.M)
            if strip(s) == strip(o):
                warns.append(f"code block {n} ({info}): only comments changed")
            else:
                errors.append(f"code block {n} ({info}): code changed")

    si, oi = inline_code(sbody), inline_code(obody)
    if si != oi:
        errors.append(f"inline code changed: lost {sorted((si - oi).elements())}, new {sorted((oi - si).elements())}")

    sl, ol = link_targets(sbody), link_targets(obody)
    if sl != ol:
        errors.append(f"link or image targets changed: lost {sorted((sl - ol).elements())}, new {sorted((ol - sl).elements())}")
    if link_titles(sbody) == link_titles(obody) and link_titles(sbody):
        warns.append("link titles were not translated")

    if ref_definitions(sbody) != ref_definitions(obody):
        errors.append(f"reference definitions changed: {ref_definitions(sbody)} -> {ref_definitions(obody)}")
    if ref_uses(sbody) != ref_uses(obody):
        errors.append(f"reference link labels changed: {dict(ref_uses(sbody))} -> {dict(ref_uses(obody))}")

    sh, oh = headings(sbody), headings(obody)
    if [l for l, _ in sh] != [l for l, _ in oh]:
        errors.append(f"heading levels changed: {[l for l, _ in sh]} -> {[l for l, _ in oh]}")
    ids = anchors(obody)
    for target in re.findall(r"\]\(#([^)\s]+)\)", obody):
        if target not in ids:
            errors.append(f"in-page link #{target} matches no heading in the translation")

    st, ot = tags(sbody), tags(obody)
    if [t for t, _ in st] != [t for t, _ in ot]:
        errors.append(f"JSX/HTML tags changed: {[t for t, _ in st]} -> {[t for t, _ in ot]}")
    else:
        for (name, sa), (_, oa) in zip(st, ot):
            if sorted(sa) != sorted(oa):
                errors.append(f"<{name}> attributes changed: {sorted(sa)} -> {sorted(oa)}")
            for k in sa:
                if k in oa and sa[k] != oa[k]:
                    warns.append(f"<{name}> {k}=\"{sa[k]}\" -> \"{oa[k]}\"")

    sc, oc = mdx_comments(sbody), mdx_comments(obody)
    if len(sc) != len(oc):
        errors.append(f"MDX comments: {len(sc)} -> {len(oc)}")
    elif sc != oc:
        warns.append("MDX comment text changed")

    s_adm = re.findall(r"^:::\w*$", sbody, re.M)
    o_adm = re.findall(r"^:::\w*$", obody, re.M)
    if s_adm != o_adm:
        errors.append(f"admonition markers changed: {s_adm} -> {o_adm}")

    s_fn, o_fn = Counter(re.findall(r"\[\^[^\]]+\]", sbody)), Counter(re.findall(r"\[\^[^\]]+\]", obody))
    if s_fn != o_fn:
        errors.append(f"footnote labels changed: {dict(s_fn)} -> {dict(o_fn)}")

    if tables(prose(sbody)) != tables(prose(obody)):
        errors.append(f"table shape changed: {tables(prose(sbody))} -> {tables(prose(obody))}")

    s_task = re.findall(r"^\s*[-*] \[([ x])\]", sbody, re.M)
    o_task = re.findall(r"^\s*[-*] \[([ x])\]", obody, re.M)
    if s_task != o_task:
        errors.append(f"task list boxes changed: {s_task} -> {o_task}")

    if hard_breaks(prose(obody)) < hard_breaks(prose(sbody)):
        errors.append(f"hard line breaks lost: {hard_breaks(prose(sbody))} -> {hard_breaks(prose(obody))}")

    for c in unmatched_emphasis(prose(obody)):
        if c not in unmatched_emphasis(prose(sbody)):
            errors.append(f"emphasis will not open or close, shows literal asterisks: ...{c}...")

    for p in tildes(prose(obody)):
        if p not in tildes(prose(sbody)):
            errors.append(f"two single ~ in one paragraph render as strikethrough in GFM: {p}")

    s_par = {p.strip() for p in re.split(r"\n\s*\n", prose(sbody)) if len(p.strip()) > 40}
    for p in re.split(r"\n\s*\n", prose(obody)):
        if p.strip() in s_par and not p.strip().startswith(("|", "<", "{", "[", "import ", "export ")):
            warns.append(f"paragraph left in English: {p.strip()[:50]}")

    for term in keep:
        if src_text.count(term) != out_text.count(term):
            warns.append(f"'{term}' appears {src_text.count(term)} times in source, {out_text.count(term)} in translation")

    return errors, warns


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("source")
    ap.add_argument("translated", nargs="+")
    ap.add_argument("--keep", action="append", default=[], help="term that must survive unchanged, e.g. a product name")
    args = ap.parse_args()
    src = open(args.source, encoding="utf-8").read()
    failed = False
    for path in args.translated:
        errors, warns = check(src, open(path, encoding="utf-8").read(), args.keep)
        print(f"== {path}: {len(errors)} errors, {len(warns)} warnings")
        for e in errors:
            print(f"ERROR   {e}")
        for w in warns:
            print(f"WARN    {w}")
        failed |= bool(errors)
    sys.exit(1 if failed else 0)


if __name__ == "__main__":
    main()