"""
tokenization_morpiece.py  —  HuggingFace wrapper around native MorPiece.

Why this exists
---------------
lm-eval loads the tokenizer through `AutoTokenizer.from_pretrained(..., trust_remote_code=True)`.
The stock HF *WordPiece* export cannot do byte fallback (only BPE/Unigram can, and
BPE merge-order does NOT reproduce MorPiece's greedy longest-match, so it would
silently perturb eng/nld). This slow `PreTrainedTokenizer` delegates every call
to the native MorPiece encoder, so:
  * byte-level fallback works identically at train time and eval time;
  * eng / nld / covered-zho segmentation is bit-identical to native MorPiece;
  * distinct rare hanzi produce distinct UTF-8 byte-token sequences -> minimal
    pairs (PinyinBench / HanziBench) stop tying -> the 0.0000 collapse is gone.

Wiring (so AutoTokenizer picks it up)
-------------------------------------
Place this file next to the model on the Hub and set, in tokenizer_config.json:
    "tokenizer_class": "MorPieceHFTokenizer",
    "auto_map": {"AutoTokenizer": ["tokenization_morpiece.MorPieceHFTokenizer", null]}
and drop the native trie next to it as `morpiece_native.json`
(that is exactly what MorPiece.save_pretrained(...) writes).
"""

import os
import json
from typing import List, Optional, Tuple

from transformers import PreTrainedTokenizer

try:
    # HuggingFace trust_remote_code copies this file into a package under
    # transformers_modules/ and imports it from there. A DOTTED-MODULE relative
    # import (`from .tokenizer_MorPiece import ...`) is the form HF follows to
    # also copy the sibling tokenizer_MorPiece.py into that package. A bare
    # `import tokenizer_MorPiece` is treated as an external PyPI dependency
    # ("Run pip install tokenizer_MorPiece"); `from . import tokenizer_MorPiece`
    # passes that check but is NOT copied, so it fails at exec time.
    from .tokenizer_MorPiece import MorPiece
except ImportError:
    # Local / same-directory use (training scripts run these as top-level
    # modules, so there is no parent package for a relative import).
    from tokenizer_MorPiece import MorPiece

NATIVE_FILE = "morpiece_native.json"


class MorPieceHFTokenizer(PreTrainedTokenizer):
    vocab_files_names = {"native_file": NATIVE_FILE}
    model_input_names = ["input_ids", "attention_mask"]

    def __init__(
        self,
        native_file: Optional[str] = None,
        unk_token="<unk>",
        pad_token="<pad>",
        bos_token="<s>",
        eos_token="</s>",
        mask_token="<mask>",
        **kwargs,
    ):
        # use_tokenizers_lib=True is REQUIRED: with it False, _preprocess_text
        # returns the RAW string (no Lowercase/NFKC, no pre-tokenisation) and
        # encode() falls back to a plain str.split(). Against a lowercase-trained
        # MorPiece vocab that byte-fragments every capital and every attached
        # punctuation mark ("The" -> <0x54> ++he, "dogs." -> dogs ++<0x2E>):
        # measured at ~21% of the token stream on ordinary eng/nld text. With it
        # True, encode() prepends a BOS symbol, which _tokenize() strips, so the
        # HF layer still owns special-token placement.
        self._mp = MorPiece(ooa=False, use_tokenizers_lib=True,
                            byte_fallback=True)
        if native_file and os.path.isfile(native_file):
            self._load_native(native_file)
        self._id_to_vocab = self._mp.id_to_vocab or {}
        self._vocab_to_id = self._mp.vocab_to_id or {}
        self._byte_set = set(getattr(self._mp, "BYTE_TOKENS", []))
        super().__init__(unk_token=unk_token, pad_token=pad_token,
                         bos_token=bos_token, eos_token=eos_token,
                         mask_token=mask_token, **kwargs)

    # -- loading -------------------------------------------------------------
    def _load_native(self, native_file: str):
        # MorPiece.from_pretrained expects a *directory* holding tokenizer.json;
        # accept either a direct file or that directory layout.
        with open(native_file, "r", encoding="utf-8") as f:
            data = json.load(f)
        self._mp.roots       = data["roots"]
        self._mp.vocab_to_id = data.get("vocab", {})
        self._mp.id_to_vocab = {v: k for k, v in self._mp.vocab_to_id.items()}
        sp = data.get("special_token_ids", {})
        self._mp.unk_token_id  = sp.get("unk",  0)
        self._mp.pad_token_id  = sp.get("pad",  1)
        self._mp.bos_token_id  = sp.get("bos",  2)
        self._mp.eos_token_id  = sp.get("eos",  3)
        self._mp.mask_token_id = sp.get("mask", 4)
        # segmentation mode: a standalone "++" glue token in the vocab means the
        # tokenizer was built with glue_morphemes (root embeddings reused,
        # "superhero" -> super ++ hero). Honour an explicit flag if present.
        self._mp.glue_morphemes = bool(
            data.get("glue_morphemes", "++" in self._mp.vocab_to_id))
        self._mp.glue_cjk_prefer_root = bool(data.get("glue_cjk_prefer_root", True))
        # Preprocessing pipeline. "native" (default) = MorPiece's own normalizer +
        # pre-tokeniser, i.e. exactly what the vocab was TRAINED with; it can emit
        # vocab entries the old cjk_safe WordPiece export could never reach
        # (apostrophe words like "don't", speaker labels like "*CHI:").
        # "legacy" = reproduce that old export bit-for-bit (Lowercase+NFKC,
        # Whitespace + .{1,24} chunker) -- use ONLY to stay id-compatible with a
        # model already trained through the old exported tokenizer.
        self._mp._pipeline = data.get("pipeline", "native")
        if self._mp._pipeline == "legacy":
            from tokenizers import normalizers, pre_tokenizers, Regex
            self._mp.normalizer = normalizers.Sequence([
                normalizers.Lowercase(), normalizers.NFKC()])
            self._mp.pre_tokenizer = pre_tokenizers.Sequence([
                pre_tokenizers.Whitespace(),
                pre_tokenizers.Split(Regex(".{1,24}"), behavior="isolated")])
        self._mp._split_re_n = -1  # force splitter-cache rebuild

    # -- vocab ---------------------------------------------------------------
    @property
    def vocab_size(self) -> int:
        return len(self._vocab_to_id)

    def get_vocab(self) -> dict:
        return dict(self._vocab_to_id, **self.added_tokens_encoder)

    # -- core delegation -----------------------------------------------------
    def _tokenize(self, text: str) -> List[str]:
        _, tokens = self._mp.encode(text)
        # encode() may emit a leading BOS symbol only when use_tokenizers_lib is
        # on; we disabled it, but strip defensively.
        if tokens and tokens[0] == self._mp.start_of_text_symbol:
            tokens = tokens[1:]
        return tokens

    def _convert_token_to_id(self, token: str) -> int:
        return self._vocab_to_id.get(token, self._mp.unk_token_id)

    def _convert_id_to_token(self, index: int) -> str:
        return self._id_to_vocab.get(index, self.unk_token)

    def convert_tokens_to_string(self, tokens: List[str]) -> str:
        """Reconstruct text with word spacing.

        MorPiece drops whitespace at encode time, so (unlike BPE) there is no
        space-marker token to invert. Word boundaries live in the root/`++`
        distinction. Two continuation encodings are supported:
          * legacy: word-internal pieces are `++X` tokens (attach, strip `++`);
          * glue_morphemes: a standalone `++` glue token precedes a root piece
            that reuses its root embedding ("superhero" -> super, ++, hero).
        A root token gets a leading space unless it is the first piece, follows
        a `++` glue, or follows a byte run. Byte-fallback runs (`<0xHH>`) fuse
        into the current word (also correct for CJK).
        """
        out: List[str] = []
        glue = False
        i, n = 0, len(tokens)
        while i < n:
            tok = tokens[i]
            if tok == "++":                                 # standalone glue
                glue = True; i += 1; continue
            if tok in self._byte_set:                       # fuse a byte run
                buf = []
                while i < n and tokens[i] in self._byte_set:
                    buf.append(int(tokens[i][3:5], 16)); i += 1
                out.append(bytes(buf).decode("utf-8", errors="replace"))
                glue = False; continue
            if tok.startswith("++"):                        # legacy ++X suffix
                out.append(tok[2:]); glue = False; i += 1; continue
            if out and not glue:                            # word-initial root
                out.append(" ")
            out.append(tok); glue = False; i += 1
        return "".join(out).replace(self._mp.SPACE_MARK, " ").strip()

    # -- saving --------------------------------------------------------------
    def save_vocabulary(self, save_directory: str,
                        filename_prefix: Optional[str] = None) -> Tuple[str]:
        os.makedirs(save_directory, exist_ok=True)
        prefix = (filename_prefix + "-") if filename_prefix else ""
        path = os.path.join(save_directory, prefix + NATIVE_FILE)
        with open(path, "w", encoding="utf-8") as f:
            json.dump({
                "roots": self._mp.roots,
                "vocab": self._vocab_to_id,
                "glue_morphemes": bool(getattr(self._mp, "glue_morphemes", False)),
                "glue_cjk_prefer_root": bool(getattr(self._mp, "glue_cjk_prefer_root", True)),
                "pipeline": getattr(self._mp, "_pipeline", "native"),
                "special_token_ids": {
                    "unk": self._mp.unk_token_id, "pad": self._mp.pad_token_id,
                    "bos": self._mp.bos_token_id, "eos": self._mp.eos_token_id,
                    "mask": self._mp.mask_token_id,
                },
            }, f, ensure_ascii=False)
        return (path,)