import gradio as gr


def chat(message, history):
    return message


def new_chat(histories):
    return [
        [{
            'summary': 'New Chat',
            'messages': [],
        }] + histories,
        0,
        []
    ]


def select_history(history_index):
    return history_index


def delete_history(histories, history_index, current_history_index):
    if current_history_index == history_index:
        current_history_index = 0
    del histories[history_index]
    return [histories, current_history_index]


def chat_change(messages, histories, history_index):
    print(messages)
    if len(messages) == 0:
        histories[history_index]['messages'] = messages
        return histories
    if len(histories) == 0:
        histories = [{}]
    histories[history_index]['messages'] = messages
    histories[history_index]['summary'] = messages[0][0] if messages[0][0].strip() else 'New Chat'
    return histories


def change_history_index(histories, current_history_index):
    return [
        histories[current_history_index]['messages'],
        {
            'histories': histories,
            'current_history_index': current_history_index,
            'event': 'change_history_index'
        }
    ]


def change_histories(histories, current_history_index):
    return {
        'histories': histories,
        'current_history_index': current_history_index,
        'event': 'change_histories'
    }

def change_read_data(data):
    if not data:
        return [[], 0, []]
    histories = data['histories']
    current_history_index = data['current_history_index']
    event = data['event']
    if event != 'load':
        return [gr.update(), gr.update(), gr.update()]
    return [histories, current_history_index, histories[current_history_index]['messages']]

with gr.Blocks(
        theme=gr.themes.Soft(),
        css="""
        .histories_container .gap { gap: 10px; }
        .history_container { gap: 10px; }
        .history_container:hover .delete_history_button { display: inline-block; }
        .history_button { display: inline-block; white-space: nowrap; overflow: hidden; text-overflow: ellipsis; }
        .delete_history_button { display: none; width: 30px; min-width: 30px; flex-grow: 0; }
        """
) as demo:
    histories = gr.State([])
    current_history_index = gr.State(0)
    write_data = gr.JSON(visible=False)
    read_data = gr.JSON(visible=False)
    with gr.Row():
        with gr.Column(scale=1):
            new_chat_button = gr.Button("New Chat", variant='primary', size='sm')

            gr.Markdown('#### Chat history')
            with gr.Column(variant='panel', elem_classes="histories_container"):
                @gr.render(inputs=histories)
                def render_todos(history_list):
                    for index, history in enumerate(history_list):
                        with gr.Row(elem_classes='history_container'):
                            history_button = gr.Button(
                                history['summary'],
                                variant='secondary',
                                size='sm',
                                elem_classes="history_button"
                            )
                            delete_history_button = gr.Button(
                                'X',
                                variant='secondary',
                                size='sm',
                                elem_classes='delete_history_button'
                            )
                            history_button.click(
                                select_history,
                                inputs=gr.State(index),
                                outputs=current_history_index
                            )
                            delete_history_button.click(
                                delete_history,
                                inputs=[histories, gr.State(index), current_history_index],
                                outputs=[histories, current_history_index]
                            )

        with gr.Column(scale=7):
            chatbot = gr.Chatbot(render=False)
            gr.ChatInterface(
                chat,
                chatbot=chatbot
            )
            chatbot.change(
                chat_change,
                inputs=[chatbot, histories, current_history_index],
                outputs=histories,
            )

        new_chat_button.click(
            new_chat,
            inputs=histories,
            outputs=[histories, current_history_index, chatbot]
        )
        current_history_index.change(
            change_history_index,
            inputs=[histories, current_history_index],
            outputs=[chatbot, write_data],
        )
        histories.change(
            change_histories,
            inputs=[histories, current_history_index],
            outputs=write_data
        )
        write_data.change(
            None,
            inputs=write_data,
            js="""
            (data) => {
              console.log(data)
              window.localStorage.setItem('data', JSON.stringify(data))
            }
            """
        )
        read_data.change(
            change_read_data,
            inputs=read_data,
            outputs=[histories, current_history_index, chatbot]
        )
        demo.load(
            None,
            outputs=read_data,
            js="""
            () => {
              const dataStr = window.localStorage.getItem('data')
              if (dataStr) {
                const data = JSON.parse(dataStr)
                data.event = 'load'
                console.log(data)
                return data
              }
              return null
            }
            """
        )

    if __name__ == "__main__":
        demo.launch()