from functools import partial
import gradio as gr

from utils import *


demo = gr.Blocks(
    css="#gradio_center {text-align: center} #char_icon {font-size: 24px; padding-top: 0px; padding-bottom: 0px;}",
    analytics_enabled=False
).queue(concurrency_count=8, api_open=False)

with demo:
    ## State Elements
    chat_init = [(None, sample_conversation.first_message)]
    chat_history = gr.State(chat_init)
    current_situation = gr.State(sample_conversation) 
    current_choices = gr.State(value=sample_conversation.starting_choices)
    user_win_state = gr.State(False)

    ## Title
    gr.Markdown('# Chat Quest!', elem_id='gradio_center')

    with gr.Column(variant='panel'):
        ## Header
        with gr.Row(variant='default'):
            with gr.Column(min_width=25):
                character_face = gr.Button(sample_conversation.character_emoji, elem_id='char_icon')

            with gr.Column(min_width=0, scale=20):
                character_name = gr.Markdown(format_chat_header(sample_conversation.character))

        ## Chat
        with gr.Column():
            chat_output = gr.Chatbot(
                value=chat_init,
                container=False,
                bubble_full_width=False,
                show_copy_button=False,
                show_share_button=False,
                height=300
            )

            situation_goal = gr.Markdown(format_goal_messsage(sample_conversation.goal_message))

            choice_buttons = [gr.Button(e, min_width=10) for e in sample_conversation.starting_choices]
            

        with gr.Column(variant='panel', visible=False) as outcome_won:
            gr.Markdown('# You have WON!', elem_id='gradio_center')
            continue_button = gr.Button('Continue')

        with gr.Column(variant='panel', visible=False) as outcome_lost:
            gr.Markdown('# You have LOST!', elem_id='gradio_center')
            retry_button = gr.Button('Retry')

    ## UI LOGIC
    for i, c in enumerate(choice_buttons):
        c.click(
            partial(update_chat, i),
            inputs=[current_choices, chat_history],
            outputs=[chat_output, chat_history] + choice_buttons,
            queue=False
        ).then(
            process_last_msg,
            inputs=[chat_history, current_situation],
            outputs=[user_win_state, current_choices, chat_output, chat_history] + choice_buttons + [outcome_won, outcome_lost]
        )

    retry_button.click(
        retry_level,
        inputs=[current_situation],
        outputs=choice_buttons + [current_choices, outcome_lost, chat_output, chat_history],
        queue=False
    )

    continue_button.click(
        load_next_level,
        inputs=[],
        outputs=choice_buttons + [current_situation, current_choices, outcome_won, chat_output, chat_history] + [character_face, character_name, situation_goal, user_win_state],
        queue=False
    )


if __name__ == "__main__":
    demo.launch(
        server_port=None,
        auth=validate_auth,
        max_threads=4, show_api=False, share=False
    )