webmail-client/backend/app/translation_prep.py

"""Getting a message ready for translation, and tidying what comes back.

A message is translated as plain text, piece by piece (the text between
tags), and the models stumble on things a reader would not: a sentence
wrapped onto two lines becomes two sentences ("ask you about" / "Plan for next
month"), "Best," becomes "Najlepsze," and a company name gets translated word
by word. So before a message goes out:

- lines that carry on a sentence (no end punctuation, next starts lowercase)
  are joined, in HTML (<br>, one <div> per line) and in wrapped plain text;
- greetings and sign-offs that make a whole line are put in the target
  language from a small phrase table (English, Polish, German);
- names are marked translate="no" and not sent — the people on the message,
  the user's own names and signature names, and capitalised multi-word names
  mid-sentence — but only the names the model would change ("Milan" →
  "Mediolan"): a kept name splits its sentence into pieces translated alone
  ("Dear <kept>Lukas Remis</kept>," → "Kochanie."), so a name the model
  leaves alone stays in the sentence;
- so are emoji, file names ("A_0023C017_p1CAC.mov") and text without a letter;
- the message is translated paragraph by paragraph (split_pieces), each from
  its own language when there are several; those in the target stay;
- and afterwards, a kept word gets back the spaces around it.
"""

from __future__ import annotations

import html as html_lib
import re
from collections.abc import Callable, Iterable
from dataclasses import dataclass, field

KEEP_OPEN = '<span translate="no">'
_TAG = re.compile(r"(<[^>]*>)")
_LOWER_START = "a-ząćęłńóśźżäöüßàâçéèêëîïôûùüÿñ"
_UPPER = "A-ZĄĆĘŁŃÓŚŹŻÄÖÜÀÂÇÉÈÊËÎÏÔÛÙÜŸÑ"

# --- joining wrapped lines -------------------------------------------------------

# A line break between a word (or comma) and a lowercase word carries a sentence on.
_HTML_SOFT_BREAK = re.compile(
    rf"([\w,;–-])[ \t]*(?:<br\s*/?>|</div>\s*<div[^>]*>|</p>\s*<p[^>]*>)\s*(?=[{_LOWER_START}])"
)
_TEXT_SOFT_BREAK = re.compile(rf"([^\s.!?:;])[ \t]*\n[ \t]*(?=[{_LOWER_START}])")


def join_wrapped_lines(html: str) -> str:
    html = _HTML_SOFT_BREAK.sub(r"\1 ", html)
    # Plain text (wm-plaintext): newlines inside text, outside tags.
    return "".join(
        p if p.startswith("<") else _TEXT_SOFT_BREAK.sub(r"\1 ", p) for p in _TAG.split(html)
    )


# --- greetings and sign-offs -----------------------------------------------------

# source language → the line as written (lower case) → target language → line.
_PHRASES: dict[str, dict[str, dict[str, str]]] = {
    "en": {
        "best": {"pl": "Pozdrawiam", "de": "Viele Grüße"},
        "best regards": {"pl": "Pozdrawiam serdecznie", "de": "Beste Grüße"},
        "best wishes": {"pl": "Najlepsze życzenia", "de": "Beste Grüße"},
        "kind regards": {"pl": "Z pozdrowieniami", "de": "Freundliche Grüße"},
        "warm regards": {"pl": "Serdecznie pozdrawiam", "de": "Herzliche Grüße"},
        "regards": {"pl": "Pozdrawiam", "de": "Grüße"},
        "cheers": {"pl": "Pozdrawiam", "de": "Viele Grüße"},
        "sincerely": {"pl": "Z poważaniem", "de": "Mit freundlichen Grüßen"},
        "yours sincerely": {"pl": "Z poważaniem", "de": "Mit freundlichen Grüßen"},
        "thanks": {"pl": "Dzięki", "de": "Danke"},
        "many thanks": {"pl": "Wielkie dzięki", "de": "Vielen Dank"},
        "thank you": {"pl": "Dziękuję", "de": "Vielen Dank"},
        "thanks in advance": {"pl": "Z góry dziękuję", "de": "Vielen Dank im Voraus"},
        "all the best": {"pl": "Wszystkiego dobrego", "de": "Alles Gute"},
        "talk soon": {"pl": "Do usłyszenia", "de": "Bis bald"},
        "have a great day": {"pl": "Miłego dnia", "de": "Einen schönen Tag noch"},
        "have a nice day": {"pl": "Miłego dnia", "de": "Einen schönen Tag noch"},
    },
    "pl": {
        "pozdrawiam": {"en": "Best regards", "de": "Viele Grüße"},
        "pozdrawiam serdecznie": {"en": "Warm regards", "de": "Herzliche Grüße"},
        "serdecznie pozdrawiam": {"en": "Warm regards", "de": "Herzliche Grüße"},
        "z poważaniem": {"en": "Kind regards", "de": "Mit freundlichen Grüßen"},
        "z pozdrowieniami": {"en": "Kind regards", "de": "Freundliche Grüße"},
        "dziękuję": {"en": "Thank you", "de": "Vielen Dank"},
        "dzięki": {"en": "Thanks", "de": "Danke"},
        "z góry dziękuję": {"en": "Thanks in advance", "de": "Vielen Dank im Voraus"},
        "miłego dnia": {"en": "Have a nice day", "de": "Einen schönen Tag noch"},
        "do usłyszenia": {"en": "Talk soon", "de": "Bis bald"},
    },
    "de": {
        "mit freundlichen grüßen": {"en": "Kind regards", "pl": "Z poważaniem"},
        "freundliche grüße": {"en": "Kind regards", "pl": "Z pozdrowieniami"},
        "viele grüße": {"en": "Best regards", "pl": "Pozdrawiam"},
        "beste grüße": {"en": "Best regards", "pl": "Pozdrawiam"},
        "liebe grüße": {"en": "Warm regards", "pl": "Serdecznie pozdrawiam"},
        "herzliche grüße": {"en": "Warm regards", "pl": "Serdecznie pozdrawiam"},
        "danke": {"en": "Thanks", "pl": "Dzięki"},
        "vielen dank": {"en": "Many thanks", "pl": "Bardzo dziękuję"},
        "bis bald": {"en": "Talk soon", "pl": "Do usłyszenia"},
    },
}
# A whole line: after a tag or line start, before a tag or line end.
_LINE = re.compile(r"(^|>|\n)([ \t]*)([^<>\n]{2,40}?)([ \t]*[,.!]?)([ \t]*)(?=<|\n|$)")


def _base(code: str) -> str:
    return code.lower().split("-")[0]


def translate_phrases(html: str, source: str, target: str) -> str:
    table = _PHRASES.get(_base(source), {})
    if not table:
        return html
    tgt = _base(target)

    def repl(m: re.Match[str]) -> str:
        found = table.get(" ".join(m.group(3).lower().split()), {}).get(tgt)
        if found is None:
            return m.group(0)
        return f"{m.group(1)}{m.group(2)}{KEEP_OPEN}{found}{m.group(4).strip()}</span>{m.group(5)}"

    # The pattern never crosses a tag: it stops at "<" and ">".
    return _LINE.sub(repl, html)


# --- names ------------------------------------------------------------------------

# A capitalised word; dots only inside it ("U.S"), never at its end — a name
# does not run on over a sentence's end ("NukeX. My work…").
_WORD = rf"[{_UPPER}][\w'’&-]*(?:\.[\w'’&-]+)*"
# Two or more capitalised words in a row ("Imaginary Pixels", "Trusted Partner
# Network"), after a lowercase word: a name, not the start of a sentence.
_MULTI_NAME = re.compile(
    rf"(?<=[{_LOWER_START},] )({_WORD}(?: (?:&|of|the|de|von|van|{_WORD}))* {_WORD})"
)
# Languages that capitalise every noun: capitals say nothing about names there.
_CAPITALISE_NOUNS = {"de", "lb"}


def _title_case(text: str) -> bool:
    """A Headline Like This: nearly every longer word capitalised (short ones
    — of, and, in — often are not). A sentence full of names is not one."""
    words = [w for w in re.findall(r"[^\W\d_][\w'’-]*", text) if len(w) >= 4]
    return len(words) >= 3 and sum(w[0].isupper() for w in words) / len(words) > 0.8


def find_names(html: str, source: str) -> set[str]:
    """Capitalised multi-word names in the middle of sentences (not in headings
    or Title Case lines, and not for languages that capitalise all nouns)."""
    if _base(source) in _CAPITALISE_NOUNS:
        return set()
    found: set[str] = set()
    for piece in _TAG.split(html):
        if piece.startswith("<") or _title_case(piece):
            continue
        found.update(m.group(1).rstrip(".") for m in _MULTI_NAME.finditer(piece))
    return found


_SIG_NAME = re.compile(rf"{_WORD}(?: (?:&|of|the|{_WORD}))* {_WORD}")


def find_names_in_signature(text: str) -> list[str]:
    """Capitalised multi-word names in a signature ("Imaginary Pixels"): the
    user's company or title, which a translation should leave alone."""
    return [m.group(0).rstrip(".") for m in _SIG_NAME.finditer(text)]


def keep_terms(html: str, terms: Iterable[str]) -> str:
    """Mark each of ``terms`` in ``html``'s text translate="no" (longest first;
    whole words; style blocks and tags left alone)."""
    words = sorted({t.strip() for t in terms if len(t.strip()) >= 2}, key=len, reverse=True)
    if not words:
        return html
    pattern = re.compile(r"(?<![\w>])(" + "|".join(re.escape(w) for w in words) + r")(?![\w<])")
    return _mark(html, pattern)


# Emoji and other pictographs (with their skin tones, joiners and variation
# selectors): on its own, an emoji comes back as a made-up word ("😊" →
# "Uzasadnienie").
_SYMBOLS = re.compile(
    "[\u2190-\u21ff\u2300-\u23ff\u2460-\u27bf\u2900-\u297f\u2b00-\u2bff"
    "\U0001f000-\U0001faff]"
    "[\u2190-\u21ff\u2300-\u23ff\u2460-\u27bf\u2900-\u297f\u2b00-\u2bff"
    "\U0001f000-\U0001faff\ufe0f\u200d\u20e3]*"
)
# File names and codes with underscores ("A_0023C017_260521_p1CAC.mov",
# "final_v2"): the models split them at every "_" and dot.
_FILE_NAME = re.compile(
    r"(?<![\w./-])(?:[\w-]*_[\w-]*(?:\.[\w-]+)*|[\w-]+(?:\.[\w-]+)*\.[a-z][a-z0-9]{1,4})(?![\w/-])",
    re.IGNORECASE,
)
# International phone numbers ("+48 780 549 877"): the models space out the
# "+". Only those: a kept piece splits its sentence, and plain numbers
# ("1300 - 2500 Euro") come through anyway.
_PHONE = re.compile(r"(?<![\w+])\+\d[\d ().-]{5,}\d(?!\w)")
# Only names with an extension a file would have, not every "e.g." or "U.S".
_FILE_EXT = re.compile(r"\.[a-z][a-z0-9]{1,4}$", re.IGNORECASE)


def keep_symbols(html: str) -> str:
    """Mark emoji, file names, handles ("Jacqueline_"), phone numbers and text
    without a letter in ``html`` translate="no"."""
    html = _mark(_mark(html, _SYMBOLS), _PHONE)

    def is_file(m: re.Match[str]) -> bool:
        word = m.group(0)
        if "_" in word and word.strip("_"):
            return True
        return bool(_FILE_EXT.search(word)) and any(c.isdigit() or c in "_-" for c in word)

    return _map_text(_mark(html, _FILE_NAME, keep=is_file), _keep_letterless)


_ENTITY = re.compile(r"&#?\w+;")
_EDGES = re.compile(r"^(\s*)(.*?)(\s*)$", re.DOTALL)


def _keep_letterless(text: str) -> str:
    """A piece of text with no letter or digit — an invisible byte-order mark,
    a lone symbol — kept as it is: the models make up a word for it ("\ufeff"
    → "Uzasadnienie")."""
    if any(c.isalnum() for c in _ENTITY.sub("", text)):
        return text
    lead, core, trail = _EDGES.match(text).groups()  # type: ignore[union-attr]
    if not _ENTITY.sub("", core).strip():
        return text  # only spaces and &nbsp;
    return f"{lead}{KEEP_OPEN}{core}</span>{trail}"


def _mark(
    html: str, pattern: re.Pattern[str], keep: Callable[[re.Match[str]], bool] | None = None
) -> str:
    """Wrap ``pattern``'s matches in ``html``'s text in translate="no" spans
    (where ``keep(match)`` says so)."""

    def repl(m: re.Match[str]) -> str:
        if keep is not None and not keep(m):
            return m.group(0)
        return f"{KEEP_OPEN}{m.group(0)}</span>"

    return _map_text(html, lambda text: pattern.sub(repl, text))


def _map_text(html: str, fn: Callable[[str], str]) -> str:
    """``fn`` applied to each piece of ``html``'s text; style blocks, tags and
    text already kept are left alone."""
    out: list[str] = []
    in_style = in_keep = False
    for piece in _TAG.split(html):
        if piece.startswith("<"):
            low = piece.lower()
            if low.startswith("<style"):
                in_style = True
            elif low.startswith("</style"):
                in_style = False
            elif low.startswith(KEEP_OPEN):
                in_keep = True
            elif in_keep and low.startswith("</span"):
                in_keep = False
            out.append(piece)
        elif in_style or in_keep or not piece:
            out.append(piece)
        else:
            out.append(fn(piece))
    return "".join(out)


# --- messages in more than one language ------------------------------------------

# Tags that end a paragraph (a line): the text between two of them is one
# piece of one language.
_BLOCK = re.compile(
    r"</?(?:p|div|li|ul|ol|table|thead|tbody|tfoot|tr|td|th|h[1-6]|blockquote|pre|"
    r"section|article|header|footer|center|br|hr)\b",
    re.IGNORECASE,
)


@dataclass
class Paragraph:
    """One paragraph of a message: its text (for telling its language) and
    where its translatable pieces are in ``split_pieces``' list."""

    text: str = ""
    pieces: list[int] = field(default_factory=list)


def split_pieces(html: str) -> tuple[list[str], list[Paragraph]]:
    """``html`` as a list of tags and text, and its paragraphs. Text in style
    blocks is no paragraph's; text already kept counts for the language but
    is not translated."""
    pieces = _TAG.split(html)
    paragraphs: list[Paragraph] = [Paragraph()]
    in_style = in_keep = False
    for i, piece in enumerate(pieces):
        if piece.startswith("<"):
            low = piece.lower()
            if low.startswith("<style"):
                in_style = True
            elif low.startswith("</style"):
                in_style = False
            elif low.startswith(KEEP_OPEN):
                in_keep = True
            elif in_keep and low.startswith("</span"):
                in_keep = False
            if _BLOCK.match(piece) and (paragraphs[-1].text.strip() or paragraphs[-1].pieces):
                paragraphs.append(Paragraph())
        elif piece and not in_style:
            paragraphs[-1].text += html_lib.unescape(piece)
            if not in_keep:
                paragraphs[-1].pieces.append(i)
    return pieces, [p for p in paragraphs if p.text.strip()]


# --- after translation -------------------------------------------------------------

# A kept span can lose the space beside it ("Bruton</span>i twoje").
# Only kept spans: they still carry translate="no" here (the sanitizer that
# runs next takes it off); other spans are the sender's and stay as they are.
_KEPT = re.compile(r'(<span translate="no">[^<]*</span>)')


def tidy(html: str) -> str:
    def repl(m: re.Match[str]) -> str:
        before = html[m.start() - 1] if m.start() > 0 else ""
        after = html[m.end()] if m.end() < len(html) else ""
        lead = " " if re.match(r"[\w,;:)]", before) else ""
        trail = " " if re.match(r"[\w(„\"“]", after) else ""
        return f"{lead}{m.group(1)}{trail}"

    return _KEPT.sub(repl, html)