import regex
from unicodedata import normalize
from .bpe_utils import bytes_to_unicode_special
from random import shuffle, seed
from math import ceil

def pretok_with_targeted_handling(text_data, split_regex, target_regex, target_process, nontarget_process):
    if split_regex:
        # 1st Regex: The general split regex
        split_regex = regex.compile(split_regex)
        split_regexed = split_regex.findall(text_data)
    else:
        split_regexed = [text_data]

    # 2nd Regex: Splits a string into target and non-target chunks
    final_processed_list = []
    for chunk in split_regexed:
        if not chunk:
            continue # Skip empty tokens if any

        # Split the token into target/non-target parts
        processed_sub_tokens = ""
        for match in regex.finditer(target_regex, chunk):
            # Check if the chunk is target to decide which process to use
            label = match.lastgroup
            content = match.group()
            if label == 'target':
                processed_sub_tokens+=(target_process(content))
            else:
                processed_sub_tokens+=(nontarget_process(content))
    
        final_processed_list.append(processed_sub_tokens)

    return final_processed_list

def pretok_with_targeted_handling_list(text_data, split_regex, target_regex, target_process, nontarget_process):
    if split_regex:
        # 1st Regex: The general split regex
        split_regex = regex.compile(split_regex)
        split_regexed = split_regex.findall(text_data)
    else:
        split_regexed = [text_data]
    # 2nd Regex: Splits a string into target and non-target chunks
    final_processed_list = []
    for chunk in split_regexed:
        if not chunk:
            continue # Skip empty tokens if any
        # Split the token into target/non-target parts
        for match in regex.finditer(target_regex, chunk):
            # Check if the chunk is target to decide which process to use
            label = match.lastgroup
            content = match.group()
            if label == 'target':
                final_processed_list+=target_process(content)
            else:
                final_processed_list+=nontarget_process(content)
    return final_processed_list


# String "좋아" -> String "ƴƸƯŠ"
def apply_encoding(input_string, encoding_type, mapping, offset=0, shuffle=False, offset_mapping=None, unicode_start=None):
    if "Jamo" in encoding_type:
        varlen = "varlen" in encoding_type
        final_string = ""
        for char in input_string:
            if char == " ":
                final_string += mapping[ord(char)]
                continue
            normalized = normalize("NFKD", char)
            if len(normalized)==1: # Single Jamo can be used, but if not usable need fallback
                if normalized not in mapping:
                    final_string += apply_encoding(normalized, "utf-8", mapping)
                    continue
            final_string += normalized
            if not varlen and len(normalized)==2: # Hangul syllable decomposed into 2 Jamo
                final_string += '∅'  # Placeholder for missing Jamo                    
        return final_string
    if shuffle:
        # For each character, calculate offset and map to shuffled offset
        input_string = "".join([chr(offset_mapping[ord(c) - unicode_start] + unicode_start) for c in input_string])
    if encoding_type == "nybbles":
        return apply_nybble_encoding(input_string, mapping, offset)
    encoded = input_string.encode(encoding_type)
    return "".join([mapping.get(x + offset) if x + offset in mapping else mapping.get(x) for x in encoded])

# String "ƴƸƯŠ"  -> String "좋아"
def apply_decoding(input_string, encoding_type, mapping, offset=0):
    antimap = {value: key for key, value in mapping.items()}
    return b"".join([antimap[a+offset].to_bytes() for a in input_string]).decode(encoding_type)

def apply_character_decoding(input_string, mapping):
    antimap = {value: key for key, value in mapping.items()}
    final_string = ""
    temp_string = b""
    char_mode = False
    for symbol in input_string:
        # For token in charmode
        if symbol not in antimap:
            # If we are in the normal case, resolve the current temp_string
            if not char_mode:
                final_string += temp_string.decode("utf8", errors='replace')
                temp_string = ""
            temp_string += symbol
            char_mode = True
        # For tokens in regular case
        else:
            # If we are in charmode, resolve the current temp_string
            if char_mode:
                final_string += temp_string
                temp_string = b""
            temp_string += (antimap[symbol]).to_bytes()
            char_mode = False
    # Resolve final piece
    if char_mode:
        final_string += temp_string
    else:
        final_string += temp_string.decode('utf8', errors='replace')
    
    return final_string

# String "è¿Ļæĺ¯ƴƸƯŠ"  -> String "这是좋아"
def apply_specialized_decoding(input_string, encoding_type, mapping, offset=0, shuffle=False, offset_mapping=None, unicode_start=None):
    if shuffle:
        def unshuffle_char(c):
            shuffled_offset = ord(c) - unicode_start
            original_offset = offset_mapping[shuffled_offset]
            return chr(original_offset + unicode_start)
    
    antimap = {value: key for key, value in mapping.items()}
    final_string = ""
    temp_string = b""
    special_mode = False
    for symbol in input_string:
        byte_id = antimap[symbol]
        # For token in special case
        if byte_id >= offset:
            # If we are in the normal case, resolve the current temp_string
            if not special_mode:
                final_string += temp_string.decode("utf8", errors='replace')
                temp_string = b""
            temp_string += (byte_id - offset).to_bytes()
            special_mode = True
        # For tokens in regular case
        else:
            # If we are in the special case, resolve the current temp_string
            if special_mode:
                temp_add = temp_string.decode(encoding_type, errors='replace')
                if shuffle:
                    temp_add = "".join([unshuffle_char(c) for c in temp_add])
                final_string += temp_add
                temp_string = b""
            temp_string += (byte_id).to_bytes()
            special_mode = False
    # Resolve final piece
    if special_mode:
        temp_add = temp_string.decode(encoding_type, errors='replace')
        if shuffle:
            temp_add = "".join([unshuffle_char(c) for c in temp_add])
        final_string += temp_add
    else:
        final_string += temp_string.decode('utf8', errors='replace')
    
    return final_string

def apply_nybble_decoding(input_string, mapping, offset=0):
    def decode_char(temp_string):
        res = []
        for i in range(0, len(temp_string), 4):
            try:
                n0, n1, n2, n3 = [t for t in temp_string[i:i+4]]
                codepoint = (n0 << 12) | (n1 << 8) | (n2 << 4) | n3
                res.append(chr(codepoint))
            except:
                # In case of any error (e.g., invalid character), append a replacement character
                res.append('�')
        return "".join(res)

    antimap = {value: key for key, value in mapping.items()}
    final_string = ""
    temp_string = b""
    special_mode = False
    for symbol in input_string:
        byte_id = antimap[symbol]
        # For token in special case
        if byte_id >= offset:
            # If we are in the normal case, resolve the current temp_string
            if not special_mode:
                final_string += temp_string.decode("utf8", errors='replace')
                temp_string = b""
            temp_string += (byte_id - offset).to_bytes()
            special_mode = True
        # For tokens in regular case
        else:
            # If we are in the special case, resolve the current temp_string
            if special_mode:
                final_string += decode_char(temp_string)
                temp_string = b""
            temp_string += (byte_id).to_bytes()
            special_mode = False
    # Resolve final piece
    if special_mode:
        final_string += decode_char(temp_string)
    else:
        final_string += temp_string.decode('utf8', errors='replace')
    
    return final_string

def apply_Jamo_decoding(input_string, mapping):
    antimap = {value: key for key, value in mapping.items()}
    final_string = ""
    current_type = type(input_string[0])
    current_pieces = []
    for char in input_string:
        char = antimap[char]
        if type(char) == current_type:
            current_pieces.append(char)
        else:
            # Process current pieces
            if current_type == str:
                # Jamo pieces
                syllables = "".join((current_pieces)).replace("∅", "")
                # Put it into unicode format
                final_string += normalize("NFKC", syllables)
            else:
                # Byte pieces
                byte_string = b"".join([c.to_bytes() for c in current_pieces]).decode('utf-8', errors='replace')
                final_string += byte_string
            current_type = type(char)
            current_pieces = [char]
    # Final flush
    if current_type == str:
        # Jamo pieces
        syllables = "".join((current_pieces)).replace("∅", "")
        final_string += syllables
    else:
        # Byte pieces
        byte_string = b"".join([c.to_bytes() for c in current_pieces]).decode('utf-8', errors='replace')
        final_string += byte_string
    return final_string 



def apply_nybble_encoding(input_string, mapping, offset,n_nybbles=4):
    def encode_char(c):
        o = ord(c)
        return "".join(mapping[offset + ((o >> (4 * i)) & 0xF)] for i in range(n_nybbles - 1, -1, -1))
    
    return "".join(encode_char(c) for c in input_string)

# Given a range of unicode, we shuffle the points
# For performance this can be stored as a list, a mapping of offset2newoffset
def create_shuffled_unicode_mapping(range_start, range_end, given_seed=6767):
    unicode_points = list(range(range_end - range_start + 1))
    seed(given_seed)
    shuffle(unicode_points)
    # to use, calculate offset = ord(char) - range_start, then map to unicode_points[offset] + range_start
    return unicode_points

# Given a range of unicode, provide a Customized UTF with given parameter properties
def create_CuTF_encoding(range_start, range_end, encode_length):
    target_space_len = range_end - range_start + 1
    num_indices = ceil(target_space_len**(1/encode_length))
    # Plan is to do the math only here and save the mapping for efficiency
    CuTF_mapping = {}
    for codepoint in range(range_start, range_end+1):
        offset = codepoint-range_start
        CuTF_rep = []
        for i in range(encode_length-1):
            rep_index = offset//(num_indices**(encode_length-i-1))
            CuTF_rep.append(rep_index)
            offset -= (num_indices**(encode_length-i-1))*rep_index
        CuTF_rep.append(offset)
        CuTF_mapping[codepoint] = CuTF_rep
    
    return CuTF_mapping, num_indices



def apply_CuTF_decoding(input_string, encode_length, num_index, range_start, indexing_strategy="shared"):
    def decode_target_chars(sequence):
        outpieces = ""
        for i in range(0, len(sequence), encode_length):
            piece = sequence[i:i+encode_length]
            offset = 0
            for j, byte in enumerate(piece):
                if indexing_strategy == "shared":
                    offset += (byte-256) * (num_index**(encode_length-j-1))
                elif indexing_strategy == "unique":
                    offset += (byte - (256 + num_index*(encode_length-j-1))) * (num_index**(encode_length-j-1))
            outpieces += chr(range_start + offset)
        return outpieces
    final_string = ""
    temp_pieces = []
    is_target = False
    for token in input_string:
        if token >= 256:
            if not is_target:
                # Flush non-target tokens
                final_string += bytes(temp_pieces).decode("utf8", errors='replace')
                temp_pieces = []
            temp_pieces.append(token)
            is_target = True
        else:
            if is_target:
                # Flush target tokens
                final_string += decode_target_chars(temp_pieces)
                temp_pieces = []
            temp_pieces.append(token)
            is_target = False
    # Flush any remaining tokens
    if is_target:
        final_string += decode_target_chars(temp_pieces)
    else:
        final_string += bytes(temp_pieces).decode("utf8", errors='replace')
    return final_string