import datetime
import gzip
import hashlib
import os
import glob
import textwrap
import gradio as gr
import numpy as np
import requests
import json
import base64
import webbrowser
from io import BytesIO
from PIL import Image
from PIL import PngImagePlugin
import random
import argparse

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

LANG = LANG_EN

# JavaScript
JAVA_SCRIPT = """
function refresh() {
    const url = new URL(window.location);

    if (url.searchParams.get('__theme') !== 'dark') {
        url.searchParams.set('__theme', 'dark');
        window.location.href = url.href;
    }
}
"""

# CSS
CSS_SCRIPT = """
#custom_prompt_text textarea {color: darkorange}
#positive_prompt_text textarea {color: greenyellow}
#negative_prompt_text textarea {color: red}
#ai_prompt_text textarea {color: hotpink}
#prompt_ban_text textarea {color: Khaki}
"""

TITLE = "WAI Character Select Preview"
CAT = "WAI_Character_Select"
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')

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

wai_illustrious_character_select_files = [
    {'name': 'wai_character', 'file_path': os.path.join(json_folder, 'wai_characters.csv'), 'url':'https://raw.githubusercontent.com/mirabarukaso/character_select_stand_alone_app/refs/heads/main/json/wai_characters.csv'},
    {'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.json'},
]
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 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_jsons():
    global character_list
    global character_dict
    global wai_image_dict
    global character_list_cn
    
    # 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:
                lines = file.readlines()
                for line in lines:
                    key, value = line.split(',')
                    character_dict[key.strip()]=value.strip()
            elif 'wai_image' == name:
                wai_image_dict = json.load(file)                
                                        
    # Create 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")

def illustrious_character_select_ex(character = 'random', optimise_tags = True, random_action_seed = 1, use_cn=False):
    chara = ''
    rnd_character = ''
    
    if 'none' == character:
        return '', '', None
    
    if 'random' == character:
        index = random_action_seed % len(character_list)
        rnd_character = character_list[index]
        if 'random' == rnd_character:
            rnd_character = character_list[index+2]
        elif 'none' == rnd_character:
            rnd_character = character_list[index+1]
    else:
        rnd_character = character
    if not use_cn:
        chara = rnd_character
    else:
        chara = character_dict[rnd_character]
        
    md5_chara = get_md5_hash(chara.replace('(','\\(').replace(')','\\)'))                
    thumb_image = Image.new('RGB', (128, 128), (128, 128, 128))
    if wai_image_dict.keys().__contains__(md5_chara):
        thumb_image = base64_to_image(wai_image_dict.get(md5_chara))
    
    opt_chara = chara
    if optimise_tags:
        opt_chara = opt_chara.replace('(', '\\(').replace(')', '\\)')  
        #print(f'{CAT}:Optimise Tags:[{chara}]->[{opt_chara}]')
    
    if not opt_chara.endswith(','):
        opt_chara = f'{opt_chara},'   
            
    return rnd_character, opt_chara, thumb_image

def refresh_character_thumb_image(character1, character2, character3):
    thumb_image = []
    if 'none' != character1 and 'random' != character1:
        _, _, thumb_image1 = illustrious_character_select_ex(character = character1)        
        thumb_image.append(thumb_image1)
    
    if 'none' != character2 and 'random' != character2:
        _, _, thumb_image2 = illustrious_character_select_ex(character = character2)        
        thumb_image.append(thumb_image2)
    if 'none' != character3 and 'random' != character3:
        _, _, thumb_image3 = illustrious_character_select_ex(character = character3, use_cn=True)        
        thumb_image.append(thumb_image3)
    return thumb_image

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