"""Emoji vocabulary: one token per complete emoji.

Every emoji in the Unicode standard is a single atomic token — single codepoints,
ZWJ sequences, flags, skin-tone variants, everything.  The vocabulary is derived
from the `emoji` library's EMOJI_DATA so it covers all ~5,225 valid emoji without
being restricted to what appears in training data.

Special tokens: PAD=0, BOS=1, EOS=2.
Emoji tokens:   indices 3..VOCAB_SIZE-1.
"""

import emoji as _emoji_lib

PAD_ID = 0
BOS_ID = 1
EOS_ID = 2
_SPECIAL = 3

# All complete emoji strings, sorted for determinism across runs
ALL_EMOJI: list[str] = sorted(_emoji_lib.EMOJI_DATA.keys())

_emoji_to_id: dict[str, int] = {em: i + _SPECIAL for i, em in enumerate(ALL_EMOJI)}
_id_to_emoji: dict[int, str] = {i + _SPECIAL: em for i, em in enumerate(ALL_EMOJI)}
VOCAB_SIZE: int = _SPECIAL + len(ALL_EMOJI)


def get_description(emoji_str: str) -> str:
    """Return a clean English description for an emoji, e.g. '🍕' → 'pizza'."""
    name = _emoji_lib.EMOJI_DATA.get(emoji_str, {}).get("en", "")
    return name.strip(":").replace("_", " ").strip()


def encode(text: str) -> list[int]:
    """Encode an emoji string → [BOS, emoji_ids..., EOS].

    Splits text into individual complete emoji using the emoji library.
    Non-emoji characters and unknown sequences are silently skipped.
    """
    tokens = [
        t.chars
        for t in _emoji_lib.analyze(text, non_emoji=False)
        if _emoji_lib.is_emoji(t.chars) and t.chars in _emoji_to_id
    ]
    return [BOS_ID] + [_emoji_to_id[em] for em in tokens] + [EOS_ID]


def decode(ids: list[int]) -> str:
    """Decode token IDs → emoji string, stripping BOS/EOS/PAD."""
    return "".join(_id_to_emoji[i] for i in ids if i in _id_to_emoji)