import gradio as gr
import sys
sys.path.append("scripts/")
from lib import init, refresh_character_thumb_image
from lib import JAVA_SCRIPT, CSS_SCRIPT, TITLE

if __name__ == '__main__':
    character_list, character_list_cn, LANG, TAG_AUTOCOMPLETE = init()
    
    def update_suggestions(text):
        matches = TAG_AUTOCOMPLETE.get_suggestions(text)
        items = []
        if matches:       
            for _, m in enumerate(matches):
                key = f"{m['prompt']} ({m['heat']})"
                items.append([key])
        return gr.Dataset(samples=items)
        
    def apply_suggestion(evt: gr.SelectData, text, custom_prompt):        
        suggestion = evt.value
        #print(f"You selected {evt.value} at {evt.index} from {evt.target} for {custom_prompt}")
        if not custom_prompt or not suggestion:
            return custom_prompt
            
        parts = custom_prompt.split(',')
        parts[-1] = suggestion[0].split(' ')[0].replace('_', ' ').replace(':', ' ')        
        return ', '.join(parts) + ', '
    
    with gr.Blocks(js=JAVA_SCRIPT, css=CSS_SCRIPT, title=TITLE) as ui:
        with gr.Row():
            character1 = gr.Dropdown(
                choices=character_list,
                label=LANG["character1"],
                value='none',
                allow_custom_value=False,
            )
            
            character2 = gr.Dropdown(
                choices=character_list,
                label=LANG["character2"],
                value='none',
                allow_custom_value=False,
            )
                            
            character3 = gr.Dropdown(
                choices=character_list_cn,
                label=LANG["character3"],
                value='none',
                allow_custom_value=False,
            )
            
        with gr.Row(elem_classes='main_row'):            
            with gr.Column(elem_classes='column_prompts'):
                    thumb_image = gr.Gallery(type="pil", columns=3, show_download_button=False, object_fit='contain', label="Thumb")    
        
        with gr.Row():
            custom_prompt = gr.Textbox(value='', label='Semi-auto tag complete test. Try tag* *tag *tag* (e621_sfw.csv@DominikDoom)') 
            suggestions = gr.Dataset(components=[custom_prompt], samples_per_page=30, samples=[])
                        
        with gr.Row():
            gr.Markdown("<h2><span style=\"color:orangered\">Thumb Image create by waiNSFWIllustrious_v120.safetensors with ComfyUI</span></h2>",)
            gr.Markdown(f"<a href='https://github.com/mirabarukaso/character_select_stand_alone_app'>Character Select SAA</a>")            
            
        character1.change(fn=refresh_character_thumb_image,
                          inputs=[character1,character2,character3],
                          outputs=[thumb_image])
        character2.change(fn=refresh_character_thumb_image,
                          inputs=[character1,character2,character3],
                          outputs=[thumb_image])
        character3.change(fn=refresh_character_thumb_image,
                          inputs=[character1,character2,character3],
                          outputs=[thumb_image])
        
        custom_prompt.change(fn=update_suggestions, inputs=[custom_prompt], outputs=[suggestions])
        suggestions.select(fn=apply_suggestion, inputs=[suggestions, custom_prompt], outputs=[custom_prompt])
        
    ui.launch()