import json
import os

import tiktoken
from transformers import PreTrainedTokenizer


def bytes_to_unicode():
    """Reversible byte -> printable-char map (GPT-2 style). Every byte maps to
    a unique character, so token strings never lose bytes."""
    bs = (list(range(ord("!"), ord("~") + 1))
          + list(range(ord("¡"), ord("¬") + 1))
          + list(range(ord("®"), ord("ÿ") + 1)))
    cs = bs[:]
    n = 0
    for b in range(256):
        if b not in bs:
            bs.append(b)
            cs.append(256 + n)
            n += 1
    return {b: chr(c) for b, c in zip(bs, cs)}


_BYTE_ENCODER = bytes_to_unicode()
_BYTE_DECODER = {v: k for k, v in _BYTE_ENCODER.items()}


class PragyaTokenizer(PreTrainedTokenizer):
    vocab_files_names = {"tokenizer_file": "pragya.tokenizer.json"}

    def __init__(self, tokenizer_file=None, bos_token="<|bos|>",
                 eos_token="<|bos|>", pad_token="<|bos|>",
                 unk_token="<|bos|>", **kwargs):
        if tokenizer_file is None:
            raise ValueError(
                "PragyaTokenizer needs tokenizer_file=<path to pragya.tokenizer.json>, "
                "or load it with from_pretrained()."
            )
        with open(tokenizer_file, "r", encoding="utf-8") as f:
            self._raw = json.load(f)

        mergeable_ranks = {bytes.fromhex(k): int(v)
                           for k, v in self._raw["mergeable_ranks"].items()}
        special_tokens = {k: int(v) for k, v in self._raw["special_tokens"].items()}
        self._enc = tiktoken.Encoding(
            name=self._raw["name"],
            pat_str=self._raw["pat_str"],
            mergeable_ranks=mergeable_ranks,
            special_tokens=special_tokens,
        )

        # id <-> token string. Ordinary tokens use the reversible byte map;
        # special tokens keep their literal string.
        self._special_set = set(special_tokens)
        self._id_to_tok = {}
        for i in range(self._enc.n_vocab):
            if i in set(special_tokens.values()):
                continue
            raw = self._enc.decode_single_token_bytes(i)
            self._id_to_tok[i] = "".join(_BYTE_ENCODER[b] for b in raw)
        for name, i in special_tokens.items():
            self._id_to_tok[i] = name
        self._tok_to_id = {t: i for i, t in self._id_to_tok.items()}

        kwargs.setdefault("clean_up_tokenization_spaces", False)
        kwargs.setdefault("padding_side", "left")  # generation needs left padding
        kwargs.setdefault("model_max_length", 1024)  # CONTEXT in config__1_.py
        super().__init__(bos_token=bos_token, eos_token=eos_token,
                         pad_token=pad_token, unk_token=unk_token, **kwargs)

    @classmethod
    def _from_pretrained(cls, resolved_vocab_files, pretrained_model_name_or_path,
                         *args, **kwargs):
        # Transformers 5 can resolve our file to a default name (tokenizer.json) that
        # we never ship. Find pragya.tokenizer.json in the local folder, or download it
        # from the Hub repo, so loading works for both paths and repo ids.
        if resolved_vocab_files.get("tokenizer_file") is None:
            name = "pragya.tokenizer.json"
            local_dir = str(pretrained_model_name_or_path)
            if os.path.isdir(local_dir) and os.path.isfile(os.path.join(local_dir, name)):
                resolved_vocab_files["tokenizer_file"] = os.path.join(local_dir, name)
            else:
                from huggingface_hub import hf_hub_download
                from huggingface_hub.errors import EntryNotFoundError
                try:
                    resolved_vocab_files["tokenizer_file"] = hf_hub_download(
                        repo_id=local_dir, filename=name,
                        revision=kwargs.get("revision"), token=kwargs.get("token"),
                    )
                except EntryNotFoundError as e:
                    raise FileNotFoundError(
                        f"{name} is missing from {local_dir} on the Hub. Upload it to the "
                        f"repository root alongside the other tokenizer files."
                    ) from e
        return super()._from_pretrained(resolved_vocab_files,
                                        pretrained_model_name_or_path,
                                        *args, **kwargs)

    def _add_tokens(self, new_tokens, special_tokens=False):
        # Transformers assigns ids past the end of the vocab to unknown tokens,
        # which the model embedding cannot look up. Refuse them instead.
        unknown = [t for t in new_tokens
                   if str(t) not in self._tok_to_id and str(t) not in self.added_tokens_encoder]
        if unknown:
            raise ValueError(
                f"Cannot add tokens {unknown!r}: this tokenizer has a fixed vocab "
                f"of {self._enc.n_vocab} tokens and the model embedding cannot grow "
                f"without retraining."
            )
        return super()._add_tokens(new_tokens, special_tokens=special_tokens)

    def tokenize(self, text, **kwargs):
        # Bypass HF's special-token splitting: the training pipeline encoded
        # "<|bos|>" in text as ordinary characters, so match that exactly.
        return self._tokenize(text, **kwargs)

    @property
    def vocab_size(self):
        return self._enc.n_vocab

    def get_vocab(self):
        return {**self._tok_to_id, **self.added_tokens_encoder}

    def _tokenize(self, text, **kwargs):
        return [self._id_to_tok[i] for i in self._enc.encode_ordinary(text)]

    def _convert_token_to_id(self, token):
        if token in self._tok_to_id:
            return self._tok_to_id[token]
        if token in self.added_tokens_encoder:
            return self.added_tokens_encoder[token]
        raise ValueError(
            f"Unknown token {token!r}. This tokenizer ships with a fixed vocab "
            f"of {self._enc.n_vocab} tokens; new special tokens cannot be added "
            f"without retraining the model embedding."
        )

    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
        # Training wrapper prepends <|bos|> by default (add_bos=True).
        ids = [self.bos_token_id] + list(token_ids_0)
        if token_ids_1 is not None:
            ids += list(token_ids_1)
        return ids

    def _convert_id_to_token(self, index):
        index = int(index)
        if index not in self._id_to_tok:
            raise ValueError(
                f"Token id {index} is outside the vocab of {self.vocab_size} tokens."
            )
        return self._id_to_tok[index]

    def _decode(self, token_ids, skip_special_tokens=False, **kwargs):
        # Accept a single id, a list, or a tensor/array (callers pass all three).
        if hasattr(token_ids, "tolist"):
            token_ids = token_ids.tolist()
        if isinstance(token_ids, int):
            token_ids = [token_ids]
        valid = [int(i) for i in token_ids if 0 <= int(i) < self.vocab_size]
        return super()._decode(valid, skip_special_tokens=skip_special_tokens, **kwargs)

    def convert_tokens_to_string(self, tokens):
        special = self._special_set | set(self.added_tokens_encoder)
        out = bytearray()
        parts = []
        for t in tokens:
            if t in special:
                parts.append(out.decode("utf-8", errors="replace"))
                out = bytearray()
                parts.append(t)
            else:
                out.extend(_BYTE_DECODER[c] for c in t)
        parts.append(out.decode("utf-8", errors="replace"))
        return "".join(parts)

    def save_vocabulary(self, save_directory, filename_prefix=None):
        name = (filename_prefix + "-" if filename_prefix else "") + "pragya.tokenizer.json"
        path = os.path.join(save_directory, name)
        with open(path, "w", encoding="utf-8") as f:
            json.dump(self._raw, f, indent=2, ensure_ascii=False)
        return (path,)