"""分词器：优先用 HuggingFace tokenizers 训 BPE，装不上就退化成字节级分词。

字节级方案零依赖、永远不会 OOV，缺点是序列变长；小语料上其实够用。
"""

import json
import os
from typing import List, Optional

try:
    from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders
    _HAS_TOKENIZERS = True
except Exception:                                    # pragma: no cover
    _HAS_TOKENIZERS = False


PAD, BOS, EOS, UNK = "<pad>", "<bos>", "<eos>", "<unk>"
SPECIALS = [PAD, BOS, EOS, UNK]


class ByteTokenizer:
    """UTF-8 字节级分词器：vocab = 4 个特殊 token + 256 个字节。"""

    kind = "byte"

    def __init__(self):
        self.vocab_size = 256 + len(SPECIALS)
        self.pad_id, self.bos_id, self.eos_id, self.unk_id = 0, 1, 2, 3
        self.offset = len(SPECIALS)

    def encode(self, text: str, bos: bool = False, eos: bool = False) -> List[int]:
        ids = [b + self.offset for b in text.encode("utf-8")]
        if bos:
            ids = [self.bos_id] + ids
        if eos:
            ids = ids + [self.eos_id]
        return ids

    def decode(self, ids: List[int]) -> str:
        buf = bytes(i - self.offset for i in ids if i >= self.offset)
        return buf.decode("utf-8", errors="replace")

    def save(self, path: str):
        with open(path, "w", encoding="utf-8") as f:
            json.dump({"kind": "byte"}, f)


class BPETokenizer:
    kind = "bpe"

    def __init__(self, tok: "Tokenizer"):
        self.tok = tok
        self.vocab_size = tok.get_vocab_size()
        self.pad_id = tok.token_to_id(PAD)
        self.bos_id = tok.token_to_id(BOS)
        self.eos_id = tok.token_to_id(EOS)
        self.unk_id = tok.token_to_id(UNK)

    def encode(self, text: str, bos: bool = False, eos: bool = False) -> List[int]:
        ids = self.tok.encode(text).ids
        if bos:
            ids = [self.bos_id] + ids
        if eos:
            ids = ids + [self.eos_id]
        return ids

    def decode(self, ids: List[int]) -> str:
        return self.tok.decode([i for i in ids if i not in (self.pad_id, self.bos_id)])

    def save(self, path: str):
        self.tok.save(path)


def train_bpe(texts: List[str], vocab_size: int, out_path: str) -> "BPETokenizer":
    if not _HAS_TOKENIZERS:
        raise RuntimeError("未安装 tokenizers 库，无法训练 BPE")
    tok = Tokenizer(models.BPE(unk_token=UNK))
    tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    tok.decoder = decoders.ByteLevel()
    trainer = trainers.BpeTrainer(vocab_size=vocab_size, special_tokens=SPECIALS,
                                  show_progress=False,
                                  initial_alphabet=pre_tokenizers.ByteLevel.alphabet())
    tok.train_from_iterator(texts, trainer=trainer)
    os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
    tok.save(out_path)
    return BPETokenizer(tok)


def load_tokenizer(path: str):
    with open(path, "r", encoding="utf-8") as f:
        head = f.read(200)
    if '"kind": "byte"' in head or '"kind":"byte"' in head:
        return ByteTokenizer()
    if not _HAS_TOKENIZERS:
        raise RuntimeError("该分词器需要 tokenizers 库")
    return BPETokenizer(Tokenizer.from_file(path))