from unittest import result

from .aether import apply_CuTF_decoding, apply_Jamo_decoding, apply_nybble_decoding, create_CuTF_encoding, pretok_with_targeted_handling, apply_encoding, apply_specialized_decoding, create_shuffled_unicode_mapping, apply_character_decoding, pretok_with_targeted_handling_list
from .consts import *
from .bpe_utils import bytes_to_unicode_original, bytes_to_unicode_special
from transformers.tokenization_utils import PreTrainedTokenizer
import os, json, shutil
from typing import Optional


# Approximately modeled after byt5 tokenizer
# https://github.com/huggingface/transformers/blob/v4.57.1/src/transformers/models/byt5/tokenization_byt5.py

class AetherByteTokenizer(PreTrainedTokenizer):
    def __init__(self, lang, encoding, shuffle=False, offset_mapping=None,**kwargs):
        self.lang = lang
        self.encoding = encoding
        self.shuffle = shuffle
        if shuffle:
            self.unicode_start = lang_ranges[lang][0]
            if offset_mapping is not None:
                self.offset_mapping = offset_mapping
            else:
                self.offset_mapping = create_shuffled_unicode_mapping(*lang_ranges[lang]) # list that acts as offset2newoffset
            self.offset_unmapping = {v: k for k, v in enumerate(self.offset_mapping)} # newoffset2offset, for decoding
        self.target_regex = target_regexes[lang]
        if encoding == "utf8":
            Aether_info = Aethers["utf8-"+lang]
            if lang == "Ko":
                # This is to limit Korean to full syllables only for this mode
                self.target_regex = target_regexes["Ko-SO"]
        else:
            Aether_info = Aethers[encoding]
        self.aether_type = Aether_info['type']
        # Mapping == aether2symbol
        self.mapping = bytes_to_unicode_special(Aether_info['pieces'], OFFSET) if self.aether_type == "bytes" else {a: a for a in Aether_info['pieces']}|bytes_to_unicode_original()
        self.index2aether = self._make_index2aether(self.mapping)
        self.index2symbol = {index: self.mapping[aether] for index, aether in self.index2aether.items()}
        self.symbol2index = {value: key for key, value in self.index2symbol.items()}
        self.target_process = lambda x: apply_encoding(x, encoding, self.mapping, OFFSET, self.shuffle, self.offset_mapping if self.shuffle else None, self.unicode_start if self.shuffle else None) #text2symbols
        self.nontarget_process = lambda x: apply_encoding(x, "utf8", self.mapping) #text2symbols

        # We want bos and eos tokens for our models
        kwargs.setdefault('bos_token', '<s>')
        kwargs.setdefault('eos_token', '</s>')
        kwargs.setdefault('special_tokens_pattern', 'bos')
        super().__init__(lang=lang, encoding=encoding, shuffle=shuffle, **kwargs)
    
    def _make_index2aether(self, mapping):
        index2aether = {index: aether_id for index, aether_id in enumerate(mapping)}
        return index2aether
        
    @property
    def vocab_size(self):
        return len(self.mapping) + len(self.added_tokens_encoder)
    
    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
        mapping_path = os.path.join(pretrained_model_name_or_path, "offset_mapping.json")
        if os.path.exists(mapping_path):
            with open(mapping_path, "r") as f:
                kwargs["offset_mapping"] = json.load(f)
        return super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)

    def get_vocab(self):
        # vocab is symbol2tok_id
        vocab = {self.convert_ids_to_tokens(i): i for i in self.index2aether}
        vocab.update(self.added_tokens_encoder)
        return vocab
    
    def pretokenize(self, input):
        return pretok_with_targeted_handling(input, None, self.target_regex, self.target_process, self.nontarget_process)

    def _tokenize(self, text: str) -> list[str]:
        """Take as input a string and return a list of strings (tokens) for words/sub-words"""
        pretokenized = self.pretokenize(text)
        tokens = []
        for chunk in pretokenized:
            tokens+=chunk
        return tokens

    def _convert_token_to_id(self, token):
        # Symbol to id
        if token in self.added_tokens_encoder:
            return self.added_tokens_encoder[token]
        tok_id = self.symbol2index[token]
        return tok_id

    def _convert_id_to_token(self, index):
        """Converts an index (integer) in a token (str) using the vocab."""
        if index in self.added_tokens_decoder:
            return self.added_tokens_decoder[index]
        return self.index2symbol[index]
    
    def _apply_decoding(self, tokens):
        if self.aether_type == "strings":
            if self.encoding == "Jamo":
                result = apply_Jamo_decoding(tokens, self.mapping)
            elif self.encoding == "Jamo-varlen":
                result = apply_Jamo_decoding(tokens, self.mapping)
        elif self.aether_type == "bytes":
            if self.encoding == "nybbles":
                result = apply_nybble_decoding(tokens, self.mapping, OFFSET)
            else:
                result = apply_specialized_decoding(tokens, self.encoding, self.mapping, OFFSET, self.shuffle, self.offset_unmapping if self.shuffle else None, self.unicode_start if self.shuffle else None)
        return result

    def convert_tokens_to_string(self, tokens):
        special_tokens_set = set(self.all_special_tokens)
        """Converts a sequence of tokens (symbols) into a single string."""
        final_string = ""
        current_iter = []
        for token in tokens:
            if token in special_tokens_set:
                # Resolve current_iter
                if current_iter:
                    resolved = self._apply_decoding(current_iter)
                    final_string += resolved
                    current_iter = []
                final_string += token
            else:
                current_iter.append(token)
        if current_iter:
            resolved = self._apply_decoding(current_iter)
            final_string += resolved
        return final_string
        

    # Tokenizer has no vocab file
    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> tuple[str]:
        return ()

    def save_pretrained(self, save_directory, **kwargs):
        saved = super().save_pretrained(save_directory, **kwargs)
        # Copy this file and all local dependencies into the directory
        src_dir = os.path.dirname(__file__)
        for fname in ["aethertokenizers.py", "aether.py", "consts.py", "bpe_utils.py"]:
            shutil.copy(os.path.join(src_dir, fname), os.path.join(save_directory, fname))
        
        if self.shuffle:
            mapping_path = os.path.join(save_directory, "offset_mapping.json")
            with open(mapping_path, "w") as f:
                json.dump(self.offset_mapping, f)
        
        config_path = os.path.join(save_directory, "tokenizer_config.json")
        with open(config_path, "r") as f:
            config = json.load(f)
        config["auto_map"] = {
            "AutoTokenizer": ["aethertokenizers.AetherByteTokenizer", None]
        }
        
        config["shuffle"] = self.shuffle

        with open(config_path, "w") as f:
            json.dump(config, f, indent=2)

        return saved

class AetherCharTokenizer(PreTrainedTokenizer):
    def __init__(self, lang, **kwargs):
        self.lang = lang
        self.unicode_start = lang_ranges[lang][0]
        self.unicode_end = lang_ranges[lang][1]
        self.target_regex = target_regexes[lang]
        if lang == "Ko":
            # This is to limit Korean to full syllables only for this mode
            self.target_regex = target_regexes["Ko-SO"]
        self.target_process = lambda x: x
        self.mapping = bytes_to_unicode_original()
        self.symbol2index = {value: key for key, value in self.mapping.items()}
        self.nontarget_process = lambda x: apply_encoding(x, "utf8", self.mapping) #text2symbols

        # We want bos and eos tokens for our models
        kwargs.setdefault('bos_token', '<s>')
        kwargs.setdefault('eos_token', '</s>')
        kwargs.setdefault('special_tokens_pattern', 'bos')
        super().__init__(lang=lang, **kwargs)
        
    @property
    def vocab_size(self):
        return len(self.mapping) + (self.unicode_end - self.unicode_start + 1) +  len(self.added_tokens_encoder)
    

    def get_vocab(self):
        # vocab is symbol2tok_id
        vocab = {self.convert_ids_to_tokens(i): i for i in self.mapping}
        vocab.update({chr(i): i - self.unicode_start + len(self.mapping) for i in range(self.unicode_start, self.unicode_end+1)})
        vocab.update(self.added_tokens_encoder)
        return vocab

    def pretokenize(self, input):
        return pretok_with_targeted_handling(input, None, self.target_regex, self.target_process, self.nontarget_process)

    def _tokenize(self, text: str) -> list[str]:
        """Take as input a string and return a list of strings (tokens) for words/sub-words"""
        pretokenized = self.pretokenize(text)
        tokens = []
        for chunk in pretokenized:
            tokens+=chunk
        return tokens

    def _convert_token_to_id(self, token):
        # Symbol to token_id
        if token in self.added_tokens_encoder:
            return self.added_tokens_encoder[token]
        if token in self.symbol2index:
            return self.symbol2index[token]
        if self.unicode_start <= ord(token) <= self.unicode_end:
            return ord(token) - self.unicode_start + len(self.mapping)
        raise ValueError(f"Token {token} not in vocab")

    def _convert_id_to_token(self, index):
        """Converts an index (integer) in a token (str) using the vocab."""
        if index in self.added_tokens_decoder:
            return self.added_tokens_decoder[index]
        if index < len(self.mapping):
            return self.mapping[index]
        return chr(index - len(self.mapping) + self.unicode_start)
    
    def _apply_decoding(self, tokens):
        result = apply_character_decoding(tokens, self.mapping)
        return result

    def convert_tokens_to_string(self, tokens):
        special_tokens_set = set(self.all_special_tokens)
        """Converts a sequence of tokens (symbols) into a single string."""
        final_string = ""
        current_iter = []
        for token in tokens:
            if token in special_tokens_set:
                # Resolve current_iter
                if current_iter:
                    resolved = self._apply_decoding(current_iter)
                    final_string += resolved
                    current_iter = []
                final_string += token
            else:
                current_iter.append(token)
        if current_iter:
            resolved = self._apply_decoding(current_iter)
            final_string += resolved
        return final_string
        

    # Tokenizer has no vocab file
    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> tuple[str]:
        return ()

    def save_pretrained(self, save_directory, **kwargs):
        saved = super().save_pretrained(save_directory, **kwargs)
        # Copy this file and all local dependencies into the directory
        src_dir = os.path.dirname(__file__)
        for fname in ["aethertokenizers.py", "aether.py", "consts.py", "bpe_utils.py"]:
            shutil.copy(os.path.join(src_dir, fname), os.path.join(save_directory, fname))
        
        config_path = os.path.join(save_directory, "tokenizer_config.json")
        with open(config_path, "r") as f:
            config = json.load(f)
        config["auto_map"] = {
            "AutoTokenizer": ["aethertokenizers.AetherCharTokenizer", None]
        }
        
        with open(config_path, "w") as f:
            json.dump(config, f, indent=2)

        return saved


# A comparable tokenizer in UTF-8.
# Could just use ByT5 too, but controls for special tokens this way

class UTF8ByteTokenizer(PreTrainedTokenizer):
    def __init__(self, **kwargs):
        self.mapping = bytes_to_unicode_special([]) #AKA index2symbol
        self.symbol2index = {value: key for key, value in self.mapping.items()}
        # We want bos and eos tokens for our models
        kwargs.setdefault('bos_token', '<s>')
        kwargs.setdefault('eos_token', '</s>')
        kwargs.setdefault('special_tokens_pattern', 'bos')
        super().__init__(**kwargs)
        
    @property
    def vocab_size(self):
        return 256 + len(self.added_tokens_encoder)
    

    def get_vocab(self):
        # vocab is symbol2tok_id
        vocab = {self.convert_ids_to_tokens(i): i for i in range(256)}
        vocab.update(self.added_tokens_encoder)
        return vocab
    

    def _tokenize(self, text: str) -> list[str]:
        """Take as input a string and return a list of strings (tokens) for words/sub-words"""
        text_bytes = text.encode('utf-8')
        tokens = [self.mapping[b] for b in text_bytes]
        return tokens

    def _convert_token_to_id(self, token):
        # Symbol to id
        if token in self.added_tokens_encoder:
            return self.added_tokens_encoder[token]
        return self.symbol2index[token]

    def _convert_id_to_token(self, index):
        """Converts an index (integer) in a token (str) using the vocab."""
        if index in self.added_tokens_decoder:
            return self.added_tokens_decoder[index]
        return self.mapping[index]

    def convert_tokens_to_string(self, tokens):
        """Converts a sequence of tokens (symbols) into a single string."""
        symbol2byte = self.get_vocab()
        final_string = ""
        curr_bstring = b""
        for token in tokens:
            if token in self.all_special_tokens:
                # Resolve current bstring
                if curr_bstring:
                    string = curr_bstring.decode("utf-8", errors="replace")
                    final_string += string
                    curr_bstring = b""
                final_string += token
            else:
                curr_bstring += bytes([symbol2byte[token]])
        if curr_bstring:
            string = curr_bstring.decode("utf-8", errors="replace")
            final_string += string
        return final_string
            

    # Tokenizer has no vocab file
    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> tuple[str]:
        return ()

    def save_pretrained(self, save_directory, **kwargs):
        saved = super().save_pretrained(save_directory, **kwargs)

        src_dir = os.path.dirname(__file__)
        for fname in ["aethertokenizers.py", "aether.py", "consts.py", "bpe_utils.py"]:
            shutil.copy(os.path.join(src_dir, fname), os.path.join(save_directory, fname))

        config_path = os.path.join(save_directory, "tokenizer_config.json")
        with open(config_path, "r") as f:
            config = json.load(f)
        config["auto_map"] = {
            "AutoTokenizer": ["aethertokenizers.UTF8ByteTokenizer", None]
        }
        with open(config_path, "w") as f:
            json.dump(config, f, indent=2)

        return saved



class AetherCuTFTokenizer(PreTrainedTokenizer):
    def __init__(self, lang, encode_length, indexing_strategy="shared", **kwargs):
        self.lang = lang
        self.encode_length = encode_length
        self.indexing_strategy = indexing_strategy
        self.target_regex = target_regexes[lang]
        if lang == "Ko":
            # This is to limit Korean to full syllables only for this mode
            self.target_regex = target_regexes["Ko-SO"]
        self.unicode_start = lang_ranges[lang][0]
        self.unicode_end = lang_ranges[lang][1]
        
        raw_mapping, self.num_CuTF_indices = create_CuTF_encoding(self.unicode_start, self.unicode_end, encode_length)
        mapping = dict()
        if indexing_strategy == "shared":
            for codepoint in raw_mapping:
                mapping[codepoint] = [256+ a for a in raw_mapping[codepoint]]
        elif indexing_strategy == "unique":
            for codepoint in raw_mapping:
                mapping[codepoint] = [256+ (encode_length -1 - i) * self.num_CuTF_indices + a for i, a in enumerate(raw_mapping[codepoint])]
        self.mapping = mapping
        self.target_process = lambda x: [rep for a in x for rep in self.mapping[ord(a)]]
        self.nontarget_process = lambda x: [a for a in x.encode('utf8')]

        # We want bos and eos tokens for our models
        kwargs.setdefault('bos_token', '<s>')
        kwargs.setdefault('eos_token', '</s>')
        kwargs.setdefault('special_tokens_pattern', 'bos')
        super().__init__(lang=lang, encode_length = encode_length, indexing_strategy= indexing_strategy, **kwargs)
    
    @property
    def vocab_size(self):
        if self.indexing_strategy == "shared":
            return 256 + self.num_CuTF_indices + len(self.added_tokens_encoder)

        elif self.indexing_strategy == "unique":
            # This is the simple answer, but in reality the full bytespace isn't utilized.
            # return 256 + self.num_CuTF_indices * self.encode_length + len(self.added_tokens_encoder)
            # Need to use the highest number as reference
            max_index = self.mapping[self.unicode_end][0]
            return max_index + len(self.added_tokens_encoder) + 1
                    
    
    def get_vocab(self):
        vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size-len(self.added_tokens_encoder))}
        vocab.update(self.added_tokens_encoder)
        return vocab
    
    def _tokenize(self, text: str) -> list[str]:
        """Take as input a string and return a list of strings (tokens) for words/sub-words"""
        return pretok_with_targeted_handling_list(text, None, self.target_regex, self.target_process, self.nontarget_process)
        
    def _convert_token_to_id(self, token):
        # Symbol to id
        if token in self.added_tokens_encoder:
            return self.added_tokens_encoder[token]
        return token

    def _convert_id_to_token(self, index):
        """Converts an index (integer) in a token (str) using the vocab."""
        if index in self.added_tokens_decoder:
            return self.added_tokens_decoder[index]
        return index
    
    def _apply_decoding(self, tokens):
        result = apply_CuTF_decoding(tokens, self.encode_length, self.num_CuTF_indices, self.unicode_start, self.indexing_strategy)
        return result

    def convert_tokens_to_string(self, tokens):
        special_tokens_set = set(self.all_special_tokens)
        """Converts a sequence of tokens (symbols) into a single string."""
        final_string = ""
        current_iter = []
        for token in tokens:
            if token in special_tokens_set:
                # Resolve current_iter
                if current_iter:
                    resolved = self._apply_decoding(current_iter)
                    final_string += resolved
                    current_iter = []
                final_string += token
            else:
                current_iter.append(token)
        if current_iter:
            resolved = self._apply_decoding(current_iter)
            final_string += resolved
        return final_string
        

    # Tokenizer has no vocab file
    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> tuple[str]:
        return ()

    def save_pretrained(self, save_directory, **kwargs):
        saved = super().save_pretrained(save_directory, **kwargs)
        # Copy this file and all local dependencies into the directory
        src_dir = os.path.dirname(__file__)
        for fname in ["aethertokenizers.py", "aether.py", "consts.py", "bpe_utils.py"]:
            shutil.copy(os.path.join(src_dir, fname), os.path.join(save_directory, fname))
        
        config_path = os.path.join(save_directory, "tokenizer_config.json")
        with open(config_path, "r") as f:
            config = json.load(f)
        config["auto_map"] = {
            "AutoTokenizer": ["aethertokenizers.AetherCuTFTokenizer", None]
        }
        
        with open(config_path, "w") as f:
            json.dump(config, f, indent=2)

        return saved