"""Text front-end: Unicode NFC, tone diacritics preserved, UTF-8 byte tokens with per-token language ids.

Code-switched input uses inline tags:  "<yor>Mo fẹ́ lọ sí ọjà</yor> <eng>to buy some bread</eng>"
Untagged text takes the default language. Each byte inherits the language of the span it is in.
"""
import re
import unicodedata

LANGS = ["eng", "ibo", "hau", "yor", "pcm"]
LANG2ID = {l: i for i, l in enumerate(LANGS)}
PAD, BOS, EOS = 0, 1, 2          # byte b -> token b + 3
VOCAB = 256 + 3
_TAG = re.compile(r"<(eng|ibo|hau|yor|pcm)>(.*?)</\1>", re.S)


def normalize(t: str) -> str:
    t = unicodedata.normalize("NFC", t)
    t = t.replace("\u2019", "'").replace("\u2018", "'").replace("\u201c", '"').replace("\u201d", '"')
    return re.sub(r"\s+", " ", t).strip()


def segments(text: str, default_lang: str = "eng"):
    """-> list of (lang, normalized_text) spans, order preserved."""
    out, pos = [], 0
    for m in _TAG.finditer(text):
        pre = text[pos:m.start()]
        if pre.strip():
            out.append((default_lang, pre))
        out.append((m.group(1), m.group(2)))
        pos = m.end()
    if text[pos:].strip():
        out.append((default_lang, text[pos:]))
    # keep a single space between spans
    res = []
    for i, (l, s) in enumerate(out):
        s = normalize(s)
        if not s:
            continue
        if res:
            s = " " + s
        res.append((l, s))
    return res


def encode(text: str, default_lang: str = "eng"):
    """-> (token ids, language ids) incl. BOS/EOS. BOS takes the first span's language, EOS the last."""
    segs = segments(text, default_lang)
    toks, langs = [BOS], [LANG2ID[segs[0][0]] if segs else LANG2ID[default_lang]]
    for l, s in segs:
        b = s.encode("utf-8")
        toks += [x + 3 for x in b]
        langs += [LANG2ID[l]] * len(b)
    toks.append(EOS)
    langs.append(langs[-1])
    return toks, langs


def plain(text: str) -> str:
    """Tagged text -> plain text (for ASR references)."""
    return " ".join(s.strip() for _, s in segments(text))