import gzip
import hashlib
import os
import requests
import json
import base64

from io import BytesIO
from PIL import Image

# Language
LANG_EN = {
    "character1": "Anima EN",
    "character2": "IL EN",
    "character3": "Anima CN",
    "character4": "IL CN",
    "action": "Action list",
    "original_character": "Original Character",    
}

LANG = LANG_EN

TITLE = "WAI Character Select Preview"
CAT =       "CSP"
CAT_IL =    "IL "
CAT_ANIMA = "ANI"
ENGLISH_CHARACTER_NAME = True

current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
json_folder = os.path.join(parent_dir, 'json')

THUMB_PLACEHOLDER_SIZE = (307, 460)

character_list = ''
character_dict = {}
wai_image_dict = {}
character_list_cn = ''
il_tag_assist = {}

anima_character_list = ''
anima_character_dict = {}
wai_anima_image_dict = {}
anima_character_list_cn = ''
anima_tag_assist = {}

PROMPT_MANAGER = None

total_counter = 0
total_counter_il = 0
total_counter_anima = 0

wai_illustrious_character_select_files = [
    {'name': 'wai_character', 'file_path': os.path.join(json_folder, 'wai_characters.csv'), 'url':'https://huggingface.co/datasets/flagrantia/character_select_stand_alone_app/resolve/main/wai_characters_v160.csv?download=true'},
    {'name': 'wai_image', 'file_path': os.path.join(json_folder, 'wai_character_thumbs.json'), 'url': 'https://huggingface.co/datasets/flagrantia/character_select_stand_alone_app/resolve/main/wai_character_thumbs_v160.json?download=true'},
    {'name': 'wai_tag_assist', 'file_path': os.path.join(json_folder, 'wai_tag_assist.json'), 'url': 'https://huggingface.co/datasets/flagrantia/character_select_stand_alone_app/resolve/main/waiIllustriousSDXL_v160_tag_assist.json?download=true'},
    
    {'name': 'wai_anima_character', 'file_path': os.path.join(json_folder, 'wai_anima_characters.csv'), 'url':'https://huggingface.co/datasets/flagrantia/character_select_stand_alone_app/resolve/main/waiANIMA_v10Base10_characters.csv?download=true'},
    {'name': 'wai_anima_image', 'file_path': os.path.join(json_folder, 'wai_anima_character_thumbs.json'), 'url': 'https://huggingface.co/datasets/flagrantia/character_select_stand_alone_app/resolve/main/waiANIMA_v10Base10_thumbs.json?download=true'},
    {'name': 'wai_anima_tag_assist', 'file_path': os.path.join(json_folder, 'wai_anima_tag_assist.json'), 'url': 'https://huggingface.co/datasets/flagrantia/character_select_stand_alone_app/resolve/main/waiANIMA_v10Base10_tag_assist.json?download=true'},
    
    # {'name': 'e621_sfw', 'file_path': os.path.join(json_folder, 'e621_sfw.csv'), 'url': 'https://raw.githubusercontent.com/DominikDoom/a1111-sd-webui-tagcomplete/refs/heads/main/tags/e621_sfw.csv'},
]
def get_md5_hash(input_str):
    md5_hash = hashlib.md5()
    md5_hash.update(input_str.encode('utf-8'))
    return md5_hash.hexdigest()

def base64_to_image(base64_data):
    compressed_data = base64.b64decode(base64_data)
    webp_data = gzip.decompress(compressed_data)
    image = Image.open(BytesIO(webp_data))
    return image

def make_thumb_placeholder(size=THUMB_PLACEHOLDER_SIZE):
    return Image.new('RGB', size, (39, 39, 42))

def ensure_thumb_size(image, size=THUMB_PLACEHOLDER_SIZE):
    if image.size != size:
        return image.resize(size, Image.LANCZOS)
    return image

def download_file(url, file_path):   
    response = requests.get(url)
    response.raise_for_status() 
    print(f'[{CAT}] Downloading... {url}')
    with open(file_path, 'wb') as file:
        file.write(response.content)        

def load_text_file(file_path):
    raw_text = ''
    if os.path.exists(file_path):
        print(f'[{CAT}] Loading {file_path}')
        with open(file_path, 'r', encoding='utf-8') as js_file:
            raw_text = js_file.read()
    else:
        print(f"[{CAT}] ERROR: {file_path} file missing!!!")
        
    return raw_text
            
def load_jsons():
    global character_list
    global character_dict
    global wai_image_dict
    global character_list_cn
    
    global anima_character_list
    global anima_character_dict
    global wai_anima_image_dict
    global anima_character_list_cn
    
    global il_tag_assist
    global anima_tag_assist
    
    global PROMPT_MANAGER
    
    # download file
    for item in wai_illustrious_character_select_files:
        name = item['name']
        file_path = item['file_path']
        url = item['url']        
        
        if not os.path.exists(file_path):
            download_file(url, file_path)
            
        with open(file_path, 'r', encoding='utf-8') as file:        
            if 'wai_character' == name:
                print(f"[{CAT}] Loading wai_character...")
                lines = file.readlines()
                for line in lines:
                    key, value = line.split(',')
                    character_dict[key.strip()]=value.strip()
            elif 'wai_image' == name:
                print(f"[{CAT}] Loading wai_image...")
                wai_image_dict = json.load(file)
            elif 'wai_anima_character' == name:
                print(f"[{CAT}] Loading wai_anima_character...")
                lines = file.readlines()
                for line in lines:
                    key, value = line.split(',')
                    anima_character_dict[key.strip()]=value.strip()
            elif 'wai_anima_image' == name:
                print(f"[{CAT}] Loading wai_anima_image...")
                wai_anima_image_dict = json.load(file)
            elif 'wai_tag_assist' == name:
                print(f"[{CAT}] Loading wai_tag_assist...")
                il_tag_assist = json.load(file)
            elif 'wai_anima_tag_assist' == name:
                print(f"[{CAT}] Loading wai_anima_tag_assist...")
                anima_tag_assist = json.load(file)
                
                                        
    # Create IL list        
    character_list = list(character_dict.values())   
    character_list.insert(0, "none")
    character_list_cn = list(character_dict.keys())   
    character_list_cn.insert(0, "none")
    
    # Create Anima List
    anima_character_list = list(anima_character_dict.values())   
    anima_character_list.insert(0, "none")
    anima_character_list_cn = list(anima_character_dict.keys())   
    anima_character_list_cn.insert(0, "none")


def illustrious_character_select_ex(character = 'none', optimise_tags = True, use_cn=False):
    global total_counter
    global total_counter_il
    
    chara = ''
    tag_assist = ''
    
    if 'none' == character:
        return '', '', None
            
    if not use_cn:
        chara = character
    else:
        chara = character_dict[character]
        
    md5_chara = get_md5_hash(chara.replace('(','\\(').replace(')','\\)'))                
    thumb_image = make_thumb_placeholder()
    if wai_image_dict.keys().__contains__(md5_chara):
        thumb_image = ensure_thumb_size(base64_to_image(wai_image_dict.get(md5_chara)))
    
    opt_chara = chara
    if optimise_tags:
        opt_chara = opt_chara.replace('(', '\\(').replace(')', '\\)')  
        total_counter = total_counter + 1
        total_counter_il = total_counter_il + 1
        print(f'{CAT_IL}:{total_counter_il}/{total_counter}:[{chara}]->[{opt_chara}]')
    
    if not opt_chara.endswith(','):
        opt_chara = f'{opt_chara},'  
    
    if il_tag_assist.keys().__contains__(chara):
        tag_assist = il_tag_assist[chara]
            
    return character, opt_chara, thumb_image, tag_assist

def anima_character_select_ex(character = 'none', optimise_tags = True, use_cn=False):
    global total_counter
    global total_counter_anima
        
    chara = ''
    tag_assist = ''
    
    if 'none' == character:
        return '', '', None
            
    if not use_cn:
        chara = character
    else:
        chara = anima_character_dict[character]
        
    md5_chara = get_md5_hash(chara.replace('(','\\(').replace(')','\\)'))                
    thumb_image = make_thumb_placeholder()
    if wai_anima_image_dict.keys().__contains__(md5_chara):
        thumb_image = ensure_thumb_size(base64_to_image(wai_anima_image_dict.get(md5_chara)))
    
    opt_chara = chara
    if optimise_tags:
        opt_chara = opt_chara.replace('(', '\\(').replace(')', '\\)')  
        total_counter = total_counter + 1
        total_counter_anima = total_counter_anima + 1
        print(f'{CAT_ANIMA}:{total_counter_anima}/{total_counter}:[{chara}]->[{opt_chara}]')
    
    if not opt_chara.endswith(','):
        opt_chara = f'{opt_chara},'   
        
    if anima_tag_assist.keys().__contains__(chara):
        tag_assist = anima_tag_assist[chara]
            
    return character, opt_chara, thumb_image, tag_assist

def create_prompt_info(rnd_character1, opt_chara1, tag_assist1,
                        rnd_character2, opt_chara2, tag_assist2,
                        rnd_character3, opt_chara3, tag_assist3,
                        rnd_character4, opt_chara4, tag_assist4):
    info = ''    
    if '' != opt_chara1:
        info += f'Anima EN: {rnd_character1}\nPrompt: {opt_chara1}\nTag Assist: {tag_assist1}\n\n'
        
    if '' != opt_chara2:
        info += f'ILLUS EN: {rnd_character2}\nPrompt: {opt_chara2}\nTag Assist: {tag_assist2}\n\n'
    
    if '' != opt_chara3:
        info += f'Anima CN: {rnd_character3}\n提示词: {opt_chara3}\n标签辅助: {tag_assist3}\n\n'
    
    if '' != opt_chara4:
        info += f'ILLUS CN: {rnd_character4}\n提示词: {opt_chara4}\n标签辅助: {tag_assist4}'

    prompt = f'{opt_chara1}{opt_chara2}{opt_chara3}{opt_chara4}'

    return prompt, info

def refresh_character_thumb_image(*args):
    thumb_image = [make_thumb_placeholder() for _ in range(4)]
    rnd_character = [''] *4
    opt_chara = [''] *4
    tag_assist = [''] * 4
    
    num_inputs = len(args)
    if num_inputs == 4:
        character1, character2, character3, character4 = args
    else:
        print(f'[refresh_character_thumb_image] Ignore old/incorrect inputs, num_inputs={num_inputs}')    
        return thumb_image, 'Refresh your page'
    
    if 'none' != character1:
        rnd_character[0], opt_chara[0], thumb_image1, tag_assist[0] = anima_character_select_ex(character = character1)        
        thumb_image[0] = thumb_image1
    if 'none' != character2:
        rnd_character[1], opt_chara[1], thumb_image2, tag_assist[1] = illustrious_character_select_ex(character = character2)
        thumb_image[1] = thumb_image2
    if 'none' != character3:
        rnd_character[2], opt_chara[2], thumb_image3, tag_assist[2] = anima_character_select_ex(character = character3, use_cn=True)
        thumb_image[2] = thumb_image3
    if 'none' != character4:
        rnd_character[3], opt_chara[3], thumb_image4, tag_assist[3] = illustrious_character_select_ex(character = character4, use_cn=True)        
        thumb_image[3] = thumb_image4
    
    _, character_info = create_prompt_info(
                        rnd_character[0], opt_chara[0], tag_assist[0],
                        rnd_character[1], opt_chara[1], tag_assist[1],
                        rnd_character[2], opt_chara[2], tag_assist[2],
                        rnd_character[3], opt_chara[3], tag_assist[3])
    
    return thumb_image, character_info

def get_prompt_manager():
    return PROMPT_MANAGER

def init():    
    load_jsons()   
    print(f'[{CAT}] Starting...')
    
    return character_list, character_list_cn, LANG, anima_character_list, anima_character_list_cn