"""
GPT tokenizer (Hugging Face `transformers`-compatible).

Wraps the project's SentencePiece model and ships inside the model repo, loaded
via `trust_remote_code=True`:

    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained("<repo>", trust_remote_code=True)
    ids = tok.build_chatml_ids([{"role": "user", "content": "Hello, who are you?"}])

`build_chatml_ids` reproduces the model's training-time ChatML tokenization
EXACTLY — including the manual injection of the `<|im_start|>` / `<|im_end|>` ids
(3 / 4) instead of encoding their literal strings. That manual injection matters:
SentencePiece would prepend a leading `▁` to a string-encoded `<|im_start|>`, and
each message segment (`role + "\n"` then `content`) is encoded separately so the
token stream matches what the model was trained on. A "build one big string, then
tokenize" chat template would NOT reproduce these boundaries, so we deliberately do
not ship a Jinja `chat_template`; `build_chatml_ids` is the authoritative encoder.
"""

import os

import sentencepiece as spm
from transformers import PreTrainedTokenizer

VOCAB_FILES_NAMES = {"vocab_file": "sentencepiece.model"}


class GPTTokenizer(PreTrainedTokenizer):
    vocab_files_names = VOCAB_FILES_NAMES
    model_input_names = ["input_ids"]

    def __init__(
        self,
        vocab_file,
        im_start_id=3,
        im_end_id=4,
        unk_token="<unk>",
        bos_token="<s>",
        eos_token="</s>",
        pad_token="<unk>",
        **kwargs,
    ):
        self.vocab_file = vocab_file
        self.sp = spm.SentencePieceProcessor(model_file=vocab_file)
        self.im_start_id = im_start_id
        self.im_end_id = im_end_id
        super().__init__(
            unk_token=unk_token,
            bos_token=bos_token,
            eos_token=eos_token,
            pad_token=pad_token,
            **kwargs,
        )

    # --- vocab ---
    @property
    def vocab_size(self):
        return self.sp.vocab_size()

    def get_vocab(self):
        vocab = {self.sp.id_to_piece(i): i for i in range(self.vocab_size)}
        vocab.update(self.added_tokens_encoder)
        return vocab

    def _tokenize(self, text):
        return self.sp.encode(text, out_type=str)

    def _convert_token_to_id(self, token):
        return self.sp.piece_to_id(token)

    def _convert_id_to_token(self, index):
        return self.sp.id_to_piece(index)

    def convert_tokens_to_string(self, tokens):
        return self.sp.decode(tokens)

    def save_vocabulary(self, save_directory, filename_prefix=None):
        out = os.path.join(
            save_directory,
            (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"],
        )
        if os.path.abspath(self.vocab_file) != os.path.abspath(out):
            with open(self.vocab_file, "rb") as src, open(out, "wb") as dst:
                dst.write(src.read())
        return (out,)

    # --- ChatML (authoritative; mirrors encode_chatml_conversation) ---
    def build_chatml_ids(self, messages, add_generation_prompt=True):
        """Token ids for a ChatML conversation.

        `messages`: list of {"role": str, "content": str}. With
        `add_generation_prompt=True`, appends `<|im_start|>assistant\\n` so the
        model continues as the assistant.
        """
        tokens = []
        for m in messages:
            tokens.append(self.im_start_id)
            tokens.extend(self.sp.encode(m["role"] + "\n"))
            tokens.extend(self.sp.encode(m["content"]))
            tokens.append(self.im_end_id)
        if add_generation_prompt:
            tokens.append(self.im_start_id)
            tokens.extend(self.sp.encode("assistant\n"))
        return tokens

    def decode_response(self, token_ids):
        """Decode generated ids to text, dropping ChatML/EOS control ids."""
        drop = {self.im_start_id, self.im_end_id, self.sp.eos_id()}
        clean = [int(t) for t in token_ids if int(t) not in drop]
        return self.sp.decode(clean).strip()