# Copyright 2026 Luke Steuber
# SPDX-License-Identifier: Apache-2.0
"""Word-level tokenizer for Chime, matching the LocalType keyboard's word rules.

A word is a letter followed by letters, apostrophes or hyphens. Words are normalized by folding
the typographic apostrophe (’) to ' and lowercasing, and the vocabulary is made of those words.

Token kinds, in the order they sit in the vocabulary:

    specials   <pad> <s> </s> <unk> <num> <sym>
    punct      one token per common punctuation mark (PUNCT below); look-alike
               characters are folded onto them (’ -> ', — -> -, 。 -> .), runs of
               the same mark collapse to one, and two or more dots become …
    emoji      every emoji the keyboard can offer, each as one token,
               ZWJ sequences and variation selectors included
    words      lowercased words, most frequent first

``<num>`` stands for any run of digits (with . , : / inside, so 12:30 and 1,000
are one token), ``<sym>`` for any other character (currency signs, stray
combining marks, control characters, an emoji the vocabulary lacks), and
``<unk>`` for a word outside the vocabulary. Whitespace is never a token.

Module-level functions work without a vocabulary; ``Tokenizer`` adds ids.
``detokenize(tokenize(text))`` is the normalized form of ``text``, and
``tokenize(detokenize(tokens)) == tokens`` for any ``tokens`` that ``tokenize``
produced.
"""

from __future__ import annotations

import re
from pathlib import Path
from typing import Iterable, Iterator, Sequence

__all__ = [
    "PAD", "BOS", "EOS", "UNK", "NUM", "SYM", "SPECIALS", "PUNCT", "CANON",
    "WORD_RE", "PREFIX_RE", "EMOJI_RE", "START",
    "normalize", "normalize_word", "active_prefix", "last_completed_word",
    "is_sentence_start", "is_word", "is_emoji", "tokenize", "detokenize",
    "strip_presentation", "Tokenizer",
]

PAD, BOS, EOS, UNK, NUM, SYM = "<pad>", "<s>", "</s>", "<unk>", "<num>", "<sym>"
SPECIALS: tuple[str, ...] = (PAD, BOS, EOS, UNK, NUM, SYM)

# Punctuation that gets its own token. Everything else non-word becomes <sym>.
PUNCT: tuple[str, ...] = (
    ".", ",", "?", "!", "…", "'", '"', "-", ":", ";", "(", ")", "/", "&",
    "$", "%", "*", "+", "=", "@", "#", "~", "_", "[", "]", "<", ">", "^",
    "\\", "|", "°",
)
_PUNCT_SET = frozenset(PUNCT)

# Look-alikes folded onto the tokens above (the device folds only ’ inside words;
# outside words it does not care, so this only shapes the model's context).
CANON: dict[str, str] = {
    "’": "'", "‘": "'", "‚": "'", "`": "'", "´": "'",
    "“": '"', "”": '"', "„": '"', "«": '"', "»": '"',
    "–": "-", "—": "-", "‐": "-", "‑": "-", "−": "-", "﹣": "-", "－": "-",
    "。": ".", "．": ".", "，": ",", "、": ",", "：": ":", "；": ";",
    "？": "?", "！": "!", "（": "(", "）": ")", "¿": "?", "¡": "!",
    "‥": "…",
}

# Stored in place of a previous word at a sentence start.
START = "\u0002"

# Python's [^\W\d_] is the L categories plus the Nl/No number letters (Ⅻ, ², ½)
# that Java's \p{L} excludes; _exact_word_pieces fixes the rare non-ASCII case.
_L = r"[^\W\d_]"
_REST = r"(?:[^\W\d_]|['’-])*"  # a negated class cannot also list characters to allow
WORD_RE = re.compile(rf"{_L}{_REST}")
PREFIX_RE = re.compile(rf"{_L}{_REST}\Z")

_EMOJI_BASE = (
    r"(?:[\U0001F1E6-\U0001F1FF]{2}"                      # flag (two regional indicators)
    r"|[0-9#*]\uFE0F?\u20E3"                              # keycap
    r"|[\U0001F000-\U0001FAFF\u2600-\u27BF\u2300-\u23FF\u2B00-\u2BFF"
    r"\u3030\u303D\u3297\u3299\u00A9\u00AE\u2122\u2194-\u2199\u21A9\u21AA"
    r"\u24C2\u25AA-\u25FE\u2934\u2935])"
)
_EMOJI_MODS = r"[\uFE0E\uFE0F\U0001F3FB-\U0001F3FF\U000E0020-\U000E007F]*"
EMOJI_RE = re.compile(rf"{_EMOJI_BASE}{_EMOJI_MODS}(?:\u200D{_EMOJI_BASE}{_EMOJI_MODS})*")

_MASTER = re.compile(
    rf"(?P<w>{_L}{_REST})"
    rf"|(?P<e>{EMOJI_RE.pattern})"
    r"|(?P<n>\d+(?:[.,:/]\d+)*)"
    r"|(?P<s>\s+)"
    r"|(?P<p>.)",
    re.DOTALL,
)

# ------------------------------------------------------------------ keyboard word rules

def normalize(word: str) -> str:
    """Fold the typographic apostrophe and lowercase."""
    return word.replace("’", "'").lower()

def _is_exact_word(token: str) -> bool:
    if token.isascii():
        return True
    return token[0].isalpha() and all(c.isalpha() or c in "'’-" for c in token)

def _exact_word_pieces(token: str) -> Iterator[tuple[bool, str]]:
    """Split a regex word that contains Nl/No characters the device would not accept."""
    run = ""
    for c in token:
        if c.isalpha() or (run and c in "'’-"):
            run += c
        else:
            if run:
                yield True, run
                run = ""
            yield False, c
    if run:
        yield True, run

def normalize_word(text: str) -> str | None:
    """The first word in ``text``, normalized."""
    for m in WORD_RE.finditer(text):
        tok = m.group()
        if _is_exact_word(tok):
            return normalize(tok)
        for is_w, piece in _exact_word_pieces(tok):
            if is_w:
                return normalize(piece)
    return None

def active_prefix(before_cursor: str) -> str:
    """The raw word being typed at the cursor, or ''."""
    i = len(before_cursor)
    while i > 0 and (before_cursor[i - 1].isalpha() or before_cursor[i - 1] in "'’-"):
        i -= 1
    run = before_cursor[i:]
    for j, c in enumerate(run):  # the prefix must start with a letter
        if c.isalpha():
            return run[j:]
    return ""

def is_sentence_start(before_cursor: str) -> bool:
    """True for an empty field, or when the word follows . ! ? or a newline."""
    tail = before_cursor[: len(before_cursor) - len(active_prefix(before_cursor))].rstrip(" \t")
    return not tail or tail[-1] in ".!?\n"

def last_completed_word(before_cursor: str) -> str | None:
    """The normalized word before the one being typed.

    Punctuation, numbers and sentence ends reset context, so "Done. The" gives None.
    """
    tail = before_cursor[: len(before_cursor) - len(active_prefix(before_cursor))].rstrip(" \t")
    if not tail or not tail[-1].isalpha():
        return None
    last = None
    for m in WORD_RE.finditer(tail):
        last = m.group()
    if last is None:
        return None
    if not _is_exact_word(last):
        last = [p for is_w, p in _exact_word_pieces(last) if is_w][-1]
    return normalize(last)

# ------------------------------------------------------------------ tokenizing

def is_word(token: str) -> bool:
    """True for a word token (not special, punctuation or emoji)."""
    return bool(token) and token[0].isalpha() and token not in SPECIALS and token not in _PUNCT_SET \
        and not EMOJI_RE.fullmatch(token)

def is_emoji(token: str) -> bool:
    return bool(EMOJI_RE.fullmatch(token))

def strip_presentation(emoji: str) -> str:
    """Key for matching an emoji regardless of variation selectors (☀ and ☀️ are one)."""
    return emoji.replace("️", "").replace("︎", "")

def _emit(out: list[str], tok: str) -> None:
    prev = out[-1] if out else None
    if tok == ".":
        if prev == ".":
            out[-1] = "…"
            return
        if prev == "…":
            return
    elif tok == prev and tok in _PUNCT_SET:
        return
    out.append(tok)

def tokenize(text: str) -> list[str]:
    """Split ``text`` into normalized tokens: words, punct, emoji, <num> and <sym>."""
    out: list[str] = []
    for m in _MASTER.finditer(text):
        kind = m.lastgroup
        if kind == "w":
            tok = m.group()
            if _is_exact_word(tok):
                out.append(normalize(tok))
            else:
                for is_w, piece in _exact_word_pieces(tok):
                    if is_w:
                        out.append(normalize(piece))
                    else:
                        _emit(out, CANON.get(piece, piece) if piece in _PUNCT_SET or piece in CANON else SYM)
        elif kind == "s":
            continue
        elif kind == "e":
            out.append(m.group())
        elif kind == "n":
            out.append(NUM)
        else:
            c = m.group()
            c = CANON.get(c, c)
            _emit(out, c if c in _PUNCT_SET else SYM)
    return out

_NO_SPACE_BEFORE = frozenset({".", ",", "?", "!", "…", ":", ";", ")", "%", "]"})
_NO_SPACE_AFTER = frozenset({"(", "[", "$", "#", "@"})
_RENDER = {NUM: "0", SYM: "¤"}  # both tokenize back to themselves

def detokenize(tokens: Iterable[str]) -> str:
    """Join tokens into readable text; the inverse of ``tokenize`` on its own output."""
    parts: list[str] = []
    glue = False
    for tok in tokens:
        if tok in (BOS, EOS, PAD):
            continue
        shown = _RENDER.get(tok, tok)
        if parts and not glue and tok not in _NO_SPACE_BEFORE:
            parts.append(" ")
        parts.append(shown)
        glue = tok in _NO_SPACE_AFTER
    return "".join(parts)

# ------------------------------------------------------------------ vocabulary

class Tokenizer:
    """A fixed vocabulary: one token per line of ``vocab.txt``, id = line number."""

    def __init__(self, tokens: Sequence[str]):
        tokens = list(tokens)
        if tuple(tokens[: len(SPECIALS)]) != SPECIALS:
            raise ValueError(f"vocabulary must start with {SPECIALS}")
        if len(set(tokens)) != len(tokens):
            raise ValueError("vocabulary has duplicate tokens")
        self.tokens: list[str] = tokens
        self.ids: dict[str, int] = {t: i for i, t in enumerate(tokens)}
        self.pad_id, self.bos_id, self.eos_id, self.unk_id, self.num_id, self.sym_id = range(6)
        self._emoji_by_key: dict[str, str] = {}
        for t in tokens:
            if is_emoji(t):
                self._emoji_by_key.setdefault(strip_presentation(t), t)
        self.word_ids: list[int] = [i for i, t in enumerate(tokens) if is_word(t)]
        self.emoji_ids: list[int] = [i for i, t in enumerate(tokens) if is_emoji(t)]
        self.punct_ids: list[int] = [i for i, t in enumerate(tokens) if t in _PUNCT_SET]

    @classmethod
    def load(cls, path: str | Path) -> "Tokenizer":
        with open(path, encoding="utf-8") as f:
            return cls([line.rstrip("\n") for line in f if line != "\n"])

    def save(self, path: str | Path) -> None:
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w", encoding="utf-8") as f:
            for t in self.tokens:
                f.write(t + "\n")

    def __len__(self) -> int:
        return len(self.tokens)

    @property
    def vocab_size(self) -> int:
        return len(self.tokens)

    def canonical(self, token: str) -> str:
        """Map an emoji typed with or without a variation selector onto the vocabulary's form."""
        if token in self.ids or not is_emoji(token):
            return token
        return self._emoji_by_key.get(strip_presentation(token), token)

    def tokens_of(self, text: str) -> list[str]:
        return [self.canonical(t) for t in tokenize(text)]

    def id_of(self, token: str) -> int:
        return self.ids.get(self.canonical(token), self.unk_id)

    def encode_tokens(self, tokens: Iterable[str], bos: bool = True, eos: bool = True) -> list[int]:
        ids = [self.id_of(t) for t in tokens]
        if bos:
            ids.insert(0, self.bos_id)
        if eos:
            ids.append(self.eos_id)
        return ids

    def encode(self, text: str, bos: bool = True, eos: bool = True) -> list[int]:
        return self.encode_tokens(tokenize(text), bos=bos, eos=eos)

    def decode(self, ids: Iterable[int]) -> list[str]:
        return [self.tokens[i] for i in ids]

    def text(self, ids: Iterable[int]) -> str:
        return detokenize(self.decode(ids))

    def is_word_id(self, i: int) -> bool:
        return is_word(self.tokens[i])