"""
GembyTiny Tokenizer — HuggingFace PreTrainedTokenizer compatible.
Works with pipeline(), Trainer, AutoTokenizer, push_to_hub().

For production: swap the character-level encoding with a SentencePiece model.
"""

import os
import json
from typing import Dict, List, Optional, Tuple
from transformers import PreTrainedTokenizer


VOCAB_FILES_NAMES = {"vocab_file": "vocab.json"}

SPECIAL_TOKENS = {
    "<|pad|>":       0,
    "<|bos|>":       1,
    "<|eos|>":       2,
    "<|sep|>":       3,
    "<|user|>":      4,
    "<|assistant|>": 5,
    "<|system|>":    6,
    "<|unk|>":       7,
}

# Default chat template (Jinja2 format — used by apply_chat_template)
CHAT_TEMPLATE = (
    "{% if messages[0]['role'] == 'system' %}"
    "<|bos|><|system|>{{ messages[0]['content'] }}<|sep|>"
    "{% set messages = messages[1:] %}"
    "{% else %}"
    "<|bos|><|system|>You are GembyTiny, a helpful AI assistant.<|sep|>"
    "{% endif %}"
    "{% for message in messages %}"
    "{% if message['role'] == 'user' %}"
    "<|user|>{{ message['content'] }}<|sep|>"
    "{% elif message['role'] == 'assistant' %}"
    "<|assistant|>{{ message['content'] }}<|sep|>"
    "{% endif %}"
    "{% endfor %}"
    "{% if add_generation_prompt %}<|assistant|>{% endif %}"
)


class GembyTinyTokenizer(PreTrainedTokenizer):
    """
    HuggingFace-compatible tokenizer for GembyTiny.

    Special tokens:
      <|pad|>       pad token
      <|bos|>       beginning of sequence
      <|eos|>       end of sequence
      <|sep|>       turn separator
      <|user|>      user turn
      <|assistant|> assistant turn
      <|system|>    system prompt
      <|unk|>       unknown token
    """

    vocab_files_names = VOCAB_FILES_NAMES
    model_input_names = ["input_ids", "attention_mask"]

    def __init__(
        self,
        vocab_file: Optional[str] = None,
        pad_token: str = "<|pad|>",
        bos_token: str = "<|bos|>",
        eos_token: str = "<|eos|>",
        unk_token: str = "<|unk|>",
        sep_token: str = "<|sep|>",
        **kwargs,
    ):
        self.vocab_file = vocab_file
        self._vocab: Dict[str, int] = {}
        self._ids_to_tokens: Dict[int, str] = {}

        if vocab_file and os.path.exists(vocab_file):
            with open(vocab_file) as f:
                self._vocab = json.load(f)
        else:
            self._build_char_vocab()

        self._ids_to_tokens = {v: k for k, v in self._vocab.items()}

        kwargs.setdefault("chat_template", CHAT_TEMPLATE)

        super().__init__(
            pad_token=pad_token,
            bos_token=bos_token,
            eos_token=eos_token,
            unk_token=unk_token,
            sep_token=sep_token,
            **kwargs,
        )

    def _build_char_vocab(self):
        """Build a character-level vocabulary (bootstrap / dev)."""
        self._vocab = dict(SPECIAL_TOKENS)
        chars = (
            list(" abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
            + list("0123456789")
            + list(".,!?;:'\"-()[]{}@#$%^&*+=/<>\\|`~\n\t")
        )
        offset = max(SPECIAL_TOKENS.values()) + 1
        for i, ch in enumerate(chars):
            self._vocab[ch] = offset + i
        self._ids_to_tokens = {v: k for k, v in self._vocab.items()}

    @property
    def vocab_size(self) -> int:
        return 32000  # fixed vocab size matching model config

    def get_vocab(self) -> Dict[str, int]:
        return dict(self._vocab)

    def _tokenize(self, text: str) -> List[str]:
        """Character-level tokenization. Replace with BPE/SPM in production."""
        return list(text)

    def _convert_token_to_id(self, token: str) -> int:
        return self._vocab.get(token, self._vocab.get("<|unk|>", 7))

    def _convert_id_to_token(self, index: int) -> str:
        return self._ids_to_tokens.get(index, "<|unk|>")

    def convert_tokens_to_string(self, tokens: List[str]) -> str:
        return "".join(tokens)

    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
        os.makedirs(save_directory, exist_ok=True)
        vocab_file = os.path.join(
            save_directory,
            (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"],
        )
        with open(vocab_file, "w", encoding="utf-8") as f:
            json.dump(self._vocab, f, ensure_ascii=False, indent=2)
        return (vocab_file,)