#!/usr/bin/env python3
"""Stellt Rohabschrift und korrigierten Text gegenüber: welche Wörter hat das Modell
ersetzt, eingefügt oder gestrichen? Ausgabe als HTML-Bericht."""
from collections import Counter
from dataclasses import dataclass, field
from difflib import SequenceMatcher
from functools import lru_cache
from pathlib import Path
import argparse
import html
import re

from vtt2txt import vtt_to_lines

# "[prüfen: …]" in German transcripts, "[check: …]" in English ones
REVIEW_MARK = re.compile(r"\[(?:prüfen|check):[^\]]*\]")
TOKEN = re.compile(rf"{REVIEW_MARK.pattern}|\S+")
NON_WORD = re.compile(r"[^\w]+")
NON_WORD_KEEP_APOSTROPHE = re.compile(r"[^\w']+")
APOSTROPHES = str.maketrans("’‘`´", "''''")
ABBREVIATION = re.compile(r"\w\.")          # one part of "z. B.", "d. h."
ABBREVIATION_END = re.compile(r"(?:^|\s)\w\.$")
EMPTY = "∅"

# System word lists (Debian/Ubuntu packages wngerman, wamerican): a replaced
# word found there was a real word, not a recognition error
WORDLISTS = {
    "de": Path("/usr/share/dict/ngerman"),
    "en": Path("/usr/share/dict/american-english"),
}
WORDLIST_PACKAGES = {"de": "wngerman", "en": "wamerican"}
MIN_SIMILARITY = 0.6       # below this a replacement counts as heavily changed
FILLERS = {"äh", "ähm", "öh", "öhm", "hm", "hmm", "uh", "uhm", "um", "er", "erm"}
CONTEXT_WORDS = 6

VALID_WORD = "gültiges Wort ersetzt"
CHANGED_A_LOT = "stark verändert"
INSERTED = "eingefügt"
DELETED = "gestrichen"


def add_token(words, token):
    """Append a token; attach punctuation-only tokens ("–") to the previous word
    and join split abbreviations ("z. B."), so neither counts as a word change."""
    if words and (
        not NON_WORD.sub("", token)
        or (ABBREVIATION.fullmatch(token) and ABBREVIATION_END.search(words[-1]))
    ):
        words[-1] += " " + token
    else:
        words.append(token)


def normalize(word):
    """Compare words without case and punctuation; those changes are expected.

    Apostrophes inside a word count ("Bahr’s" is not "Bahrs"), quotes at its
    edges do not.
    """
    word = word.lower().translate(APOSTROPHES)
    return NON_WORD_KEEP_APOSTROPHE.sub("", word).strip("'") or word


@dataclass
class Comparison:
    raw: list                     # words of the raw transcript
    words: list                   # words of the corrected text, without review marks
    paragraph: list               # paragraph number of each corrected word
    marks: dict = field(default_factory=dict)  # word index -> review marks after it
    opcodes: list = field(default_factory=list)
    suspects: list = field(default_factory=list)  # Suspect, see review_candidates

    def changes(self):
        """Yield (tag, raw words, corrected words) for every difference."""
        for tag, i1, i2, j1, j2 in self.opcodes:
            if tag != "equal":
                yield tag, self.raw[i1:i2], self.words[j1:j2]


def compare(raw, corrected):
    """Word-level comparison of the raw transcript and the corrected text."""
    words, paragraph, marks = [], [], {}
    for number, text in enumerate(re.split(r"\n\s*\n", corrected.strip())):
        start = len(words)
        for token in TOKEN.findall(text):
            # Review marks stay out of the comparison, but are shown after their word
            if REVIEW_MARK.fullmatch(token):
                marks.setdefault(len(words) - 1, []).append(token)
            # Punctuation is never merged across paragraphs
            elif len(words) == start:
                words.append(token)
            else:
                add_token(words, token)
        paragraph += [number] * (len(words) - start)

    raw_words = []
    for token in raw.split():
        add_token(raw_words, token)
    matcher = SequenceMatcher(
        None, [normalize(w) for w in raw_words], [normalize(w) for w in words],
        autojunk=False,
    )
    return Comparison(raw_words, words, paragraph, marks, matcher.get_opcodes())


@dataclass
class Suspect:
    """A change by the model worth a look; j1:j2 are corrected word indices."""
    j1: int
    j2: int
    old: str
    new: str
    reason: str


@lru_cache(maxsize=None)
def load_wordlist(lang):
    """Lower-case words of the system word list, or None if it is missing.

    Read once per run (the German list has some 350,000 words).
    """
    path = WORDLISTS.get(lang)
    if path is None or not path.is_file():
        package = WORDLIST_PACKAGES.get(lang)
        hint = f" (Paket {package})" if package else ""
        print(
            f"Hinweis: keine Wortliste für „{lang}“{hint}; "
            "ersetzte gültige Wörter werden nicht erkannt."
        )
        return None
    with path.open(encoding="utf-8", errors="replace") as words:
        return {line.strip().lower() for line in words if line.strip()}


def glossary_words(terms, replacements):
    """Normalized words of glossary terms and replacement targets."""
    texts = list(terms) + [right for _, right in replacements]
    return {normalize(word) for text in texts for word in text.split()}


def spelling_key(words):
    """Letters only: joining, splitting and hyphens make no difference."""
    return re.sub(r"[\W_]+", "", "".join(words).lower())


def review_candidates(comparison, wordlist=None, known=()):
    """Changes the model made on its own that may alter the meaning.

    Not reported: changes of spacing, hyphens and punctuation only, new words
    from the glossary, deleted filler words, and replacements of words not in
    the word list (recognition errors) that stay similar.
    """
    suspects = []
    for tag, i1, i2, j1, j2 in comparison.opcodes:
        if tag == "equal":
            continue
        old, new = comparison.raw[i1:i2], comparison.words[j1:j2]
        if new and all(normalize(word) in known for word in new):
            continue
        if tag == "delete":
            if all(normalize(word) in FILLERS for word in old):
                continue
            reason = DELETED
        elif tag == "insert":
            reason = INSERTED
        else:
            old_key, new_key = spelling_key(old), spelling_key(new)
            if old_key == new_key:
                continue
            if wordlist is not None and all(normalize(w) in wordlist for w in old):
                reason = VALID_WORD
            elif SequenceMatcher(None, old_key, new_key).ratio() < MIN_SIMILARITY:
                reason = CHANGED_A_LOT
            else:
                continue
        suspects.append(Suspect(j1, j2, " ".join(old), " ".join(new), reason))
    return suspects


def change_summary(comparison):
    """[(old, new, count)] of all changes, most frequent first."""
    counts = Counter(
        (" ".join(old), " ".join(new)) for _, old, new in comparison.changes()
    )
    return [(old, new, count) for (old, new), count in counts.most_common()]


STYLE = """
:root { --bg: #fff; --fg: #1d1d1f; --muted: #6e6e73; --rep: #fff3b0; --ins: #d9f7d9;
        --del: #fde2e2; --check: #ffe0b3; --line: #ddd; --warn: #d70015; }
@media (prefers-color-scheme: dark) {
  :root { --bg: #1c1c1e; --fg: #f2f2f7; --muted: #a1a1a6; --rep: #5c4d00; --ins: #1f4d1f;
          --del: #5c1f1f; --check: #5c3a00; --line: #3a3a3c; --warn: #ff6961; }
}
body { background: var(--bg); color: var(--fg); font: 16px/1.6 system-ui, sans-serif;
       max-width: 50rem; margin: 0 auto; padding: 1.5rem 1rem; }
h1 { font-size: 1.4rem; } h2 { font-size: 1.15rem; margin-top: 2rem; }
.muted { color: var(--muted); }
table { border-collapse: collapse; width: 100%; font-size: 0.95rem; }
td, th { border-bottom: 1px solid var(--line); padding: 0.2rem 0.5rem; text-align: left; }
td.n { text-align: right; width: 3rem; }
.rep { background: var(--rep); }
ins { background: var(--ins); text-decoration: none; }
del { background: var(--del); }
.orig { color: var(--muted); font-size: 0.8em; }
.check { background: var(--check); font-size: 0.85em; }
.warn { outline: 2px solid var(--warn); outline-offset: 1px; border-radius: 2px; }
td.ctx { color: var(--muted); } td.ctx b { color: var(--fg); }
"""


def render_html(comparison, title):
    esc = html.escape
    summary = change_summary(comparison)
    counts = Counter(tag for tag, _, _ in comparison.changes())
    review_marks = sum(len(m) for m in comparison.marks.values())

    rows = "\n".join(
        f"<tr><td class=n>{count}</td><td><del>{esc(old) or EMPTY}</del></td>"
        f"<td><ins>{esc(new) or EMPTY}</ins></td></tr>"
        for old, new, count in summary
    )
    suspects = comparison.suspects
    suspect_rows = "\n".join(
        f"<tr><td>{esc(s.reason)}</td><td><del>{esc(s.old) or EMPTY}</del></td>"
        f"<td><ins>{esc(s.new) or EMPTY}</ins></td>"
        f"<td class=ctx>{suspect_context(comparison, s)}</td></tr>"
        for s in suspects
    )
    if suspects:
        suspect_table = (
            "<table><tr><th>Grund</th><th>vorher</th><th>korrigiert</th>"
            f"<th>Zusammenhang</th></tr>\n{suspect_rows}\n</table>"
        )
    else:
        suspect_table = "<p class=muted>Keine auffälligen Änderungen gefunden.</p>"

    # Suspect words are framed; a suspect deletion shows ⚠ after the word before it
    warn_words, warn_after = {}, {}
    for s in suspects:
        note = esc(f"{s.reason}: „{s.old or EMPTY}“ → „{s.new or EMPTY}“", quote=True)
        if s.j1 == s.j2:
            warn_after.setdefault(max(s.j1 - 1, 0), []).append(note)
        for j in range(s.j1, s.j2):
            warn_words[j] = note

    # Corrected text with its paragraphs; each word knows its paragraph
    parts = []
    current = 0

    def emit(j, piece):
        nonlocal current
        if comparison.paragraph[j] != current:
            current = comparison.paragraph[j]
            parts.append("</p>\n<p>")
        if j in warn_words:
            piece = f"<span class=warn title='{warn_words[j]}'>{piece}</span>"
        parts.append(piece)
        for mark in comparison.marks.get(j, []):
            parts.append(f"<span class=check>{esc(mark)}</span>")
        for note in warn_after.get(j, []):
            parts.append(f"<span class=warn title='{note}'>⚠</span>")

    for tag, i1, i2, j1, j2 in comparison.opcodes:
        old = " ".join(comparison.raw[i1:i2])
        if tag == "delete":
            parts.append(f"<del title='gestrichen'>{esc(old)}</del>")
            continue
        for j in range(j1, j2):
            word = esc(comparison.words[j])
            if tag == "equal":
                emit(j, word)
            elif tag == "insert":
                emit(j, f"<ins title='eingefügt'>{word}</ins>")
            else:
                emit(j, f"<span class=rep title='Original: {esc(old, quote=True)}'>{word}</span>")
        if tag == "replace":
            parts.append(f"<span class=orig>({esc(old)})</span>")

    text = " ".join(parts).replace(" </p>\n<p> ", "</p>\n<p>")

    return f"""<!doctype html>
<html lang="de"><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>Änderungen: {esc(title)}</title><style>{STYLE}</style></head>
<body>
<h1>Änderungen: {esc(title)}</h1>
<p class=muted>{len(comparison.raw)} Wörter Rohtext, {len(comparison.words)} Wörter korrigiert.
{counts['replace']} ersetzt, {counts['insert']} eingefügt, {counts['delete']} gestrichen,
{review_marks} Prüfmarken, {len(suspects)} Stellen zu prüfen.
Groß-/Kleinschreibung und Satzzeichen sind nicht markiert.</p>
<p><span class=rep>ersetzt</span> <span class=orig>(Original)</span>
<ins>eingefügt</ins> <del>gestrichen</del> <span class=check>[prüfen: …]</span>
<span class=warn>zu prüfen</span></p>
<h2>Zu prüfen ({len(suspects)})</h2>
<p class=muted>Eigenmächtige Änderungen des Modells, die den Sinn ändern können:
gültige Wörter ersetzt, stark veränderte, eingefügte oder gestrichene Wörter.
„vorher“ ist der Text, den das Modell bekommen hat (nach den Glossar-Ersetzungen).</p>
{suspect_table}
<h2>Alle Änderungen nach Häufigkeit</h2>
<table><tr><th class=n>×</th><th>Rohtext</th><th>korrigiert</th></tr>
{rows}
</table>
<h2>Korrigierter Text</h2>
<p>{text}</p>
</body></html>
"""


def suspect_context(comparison, suspect):
    """HTML: the change in bold with a few corrected words around it."""
    esc = html.escape
    words = comparison.words
    before = " ".join(words[max(0, suspect.j1 - CONTEXT_WORDS):suspect.j1])
    after = " ".join(words[suspect.j2:suspect.j2 + CONTEXT_WORDS])
    if suspect.j1 == suspect.j2:
        middle = f"<del>{esc(suspect.old)}</del>"
    else:
        middle = f"<b>{esc(' '.join(words[suspect.j1:suspect.j2]))}</b>"
    return f"… {esc(before)} {middle} {esc(after)} …"


def print_summary(comparison, limit=15, suspect_limit=10):
    summary = change_summary(comparison)
    print(f"Änderungen gegenüber dem Rohtext: {sum(c for _, _, c in summary)}")
    for old, new, count in summary[:limit]:
        print(f"  {count}× „{old or EMPTY}“ → „{new or EMPTY}“")
    if len(summary) > limit:
        print(f"  … und {len(summary) - limit} weitere, siehe Bericht")

    suspects = comparison.suspects
    print(f"Zu prüfen: {len(suspects)} Stellen")
    for s in suspects[:suspect_limit]:
        print(f"  {s.reason}: „{s.old or EMPTY}“ → „{s.new or EMPTY}“")
    if len(suspects) > suspect_limit:
        print(f"  … und {len(suspects) - suspect_limit} weitere, siehe Bericht")


def write_report(raw, corrected, target, title, model_input=None, known=(), lang="de"):
    """Write the HTML report; return the comparison with the suspect changes.

    model_input is the text the model got (raw text after glossary
    replacements); only changes against it count as suspect.
    """
    comparison = compare(raw, corrected)
    by_model = compare(raw if model_input is None else model_input, corrected)
    comparison.suspects = review_candidates(by_model, load_wordlist(lang), set(known))
    target.write_text(render_html(comparison, title), encoding="utf-8")
    return comparison


def report_path(corrected_path):
    """rede.korrigiert.txt -> rede.aenderungen.html"""
    stem = corrected_path.stem.removesuffix(".korrigiert")
    return corrected_path.with_name(f"{stem}.aenderungen.html")


def read_raw(path):
    text = path.read_text(encoding="utf-8-sig")
    if path.suffix.lower() == ".vtt":
        text = "\n".join(vtt_to_lines(text))
    return text


def parse_args():
    parser = argparse.ArgumentParser(
        description="Stellt Rohabschrift und korrigierten Text als HTML-Bericht gegenüber."
    )
    parser.add_argument(
        "-a", "--roh", required=True, type=Path, help="Rohabschrift (.txt oder .vtt)"
    )
    parser.add_argument(
        "-b", "--korrigiert", required=True, type=Path, help="korrigierter Text"
    )
    parser.add_argument(
        "-o", "--output", type=Path,
        help="HTML-Bericht (Standard: <korrigiert ohne .korrigiert.txt>.aenderungen.html)",
    )
    parser.add_argument(
        "-g", "--glossar", type=Path, action="append", default=[],
        help="Glossar(e) der Korrektur: Ersetzungen und Begriffe gelten nicht als "
        "verdächtig (mehrfach angeben wie bei abschrift_korrigieren.py)",
    )
    parser.add_argument(
        "-l", "--lang", default="de",
        help="Sprache des Beitrags für die Wortliste (de, en; Standard: de)",
    )
    args = parser.parse_args()
    for path in (args.roh, args.korrigiert, *args.glossar):
        if not path.is_file():
            parser.error(f"Datei nicht gefunden: {path}")
    if args.output is None:
        args.output = report_path(args.korrigiert)
    return args


def main():
    # abschrift_korrigieren imports this module: import it only when needed
    from abschrift_korrigieren import apply_replacements, read_glossary

    args = parse_args()
    raw = read_raw(args.roh)
    terms, replacements = [], []
    for glossary in args.glossar:
        more_terms, more_replacements = read_glossary(glossary)
        terms += more_terms
        replacements += more_replacements
    # Like abschrift_korrigieren.py: replacements work on running text
    model_input = apply_replacements(" ".join(raw.split()), replacements)

    corrected = args.korrigiert.read_text(encoding="utf-8")
    comparison = write_report(
        raw, corrected, args.output, args.korrigiert.name,
        model_input=model_input, known=glossary_words(terms, replacements),
        lang=args.lang,
    )
    print_summary(comparison)
    print(f"Bericht: {args.output}")


if __name__ == "__main__":
    main()