import json
import os
from transformers import PreTrainedTokenizer


class NomTokenizer(PreTrainedTokenizer):
    vocab_files_names = {'vocab_file': 'vocab.json'}
    
    def __init__(
        self,
        vocab_file,
        unk_token='<UNK>',
        unk_token_id=0,
        id_start=1,
        **kwargs
    ):
        self.vocab_file = vocab_file
        self.id_start = id_start
        self.unk_token = unk_token
        self.unk_token_id = unk_token_id
        self.pad_token = unk_token
        self.pad_token_id = unk_token_id

        with open(vocab_file, 'r', encoding='utf-8') as f:
            self.vocab_dict = json.load(f)
    
        self.char2id = {}
        self.id2char = {}
        for i, char in enumerate(self.vocab_dict.keys(), start=id_start):
            self.char2id[char] = i
            self.id2char[i] = char
        self.id_to_options = {idx: v for idx, v in enumerate(self.vocab_dict.values(), start=id_start)}

        super().__init__(**kwargs)

    def _tokenize(self, text):
        return list(text)

    def _convert_token_to_id(self, token):
        return self.char2id.get(token, self.unk_token_id)

    def _convert_id_to_token(self, index):
        if index == self.unk_token_id:
            return self.unk_token
        return self.id2char.get(index, self.unk_token)

    @property
    def vocab_size(self):
        return len(self.char2id) + 1

    def get_vocab(self):
        vocab = {**self.char2id, **self.added_tokens_encoder}
        return vocab

    def save_vocabulary(self, save_directory, filename_prefix=None):
        if filename_prefix:
            vocab_file = os.path.join(save_directory, f'{filename_prefix}-vocab.json')
        else:
            vocab_file = os.path.join(save_directory, 'vocab.json')
        
        with open(vocab_file, 'w', encoding='utf-8') as f:
            json.dump(self.vocab_dict, f, ensure_ascii=False)
        
        return (vocab_file,)