import gradio as gr
from backend import get_handler
import queue
from db import collection
from info_table import api_info, api_samples
from app_utils import *
from state_manager import StateManager

state_manager = StateManager(collection)
shared_queue = queue.Queue()

def report_issue():
    return gr.Info("Thank you for reporting the issue. Our team will look into it.")

with gr.Blocks(css=CUSTOM_CSS) as demo:
    gr.Markdown("# Multiturn LLM Chat Interface")
    
    with gr.Tabs() as tabs:
        with gr.Tab("Single Model Demo"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("## Configuration")
                    single_model_dropdown = gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[0], interactive=True)
                    single_temperature_slider = gr.Slider(0, 1, value=DEFAULT_TEMPERATURE_1, label="Temperature", interactive=True)
                    single_category_dropdown = gr.Dropdown(choices=CATEGORIES, label="Select Category", value=CATEGORIES[0], interactive=True)
                    
                with gr.Column(scale=2):
                    single_chatbot = gr.Chatbot(elem_id="chatbot", bubble_full_width=False, type="messages")
                    with gr.Row():
                        single_restart_btn = gr.Button("Restart")
                        single_report_btn = gr.Button("Report Issue")
                    single_chat_input = gr.MultimodalTextbox(
                        interactive=True,
                        file_count="multiple",
                        placeholder="Enter message or upload file...",
                        show_label=False,
                    )
                    demo.load(lambda: get_initial_state(), outputs=single_chatbot)
                    single_chat_msg = single_chat_input.submit(
                        state_manager.add_message, [single_chat_input], [single_chatbot, single_chat_input]
                    )
                    single_bot_msg = single_chat_msg.then(state_manager.get_reponse_single, [], single_chatbot, api_name="bot_response")
                    single_bot_msg.then(lambda: gr.MultimodalTextbox(interactive=True), None, [single_chat_input])

                    single_chatbot.like(print_like_dislike, None, None, like_user_message=True)
                    
                    with gr.Row():
                        single_example_btn1 = gr.Button("Example 1 - GFSFileSystem")
                        single_example_btn2 = gr.Button("Example 2 - TradingBot")
                        single_example_btn3 = gr.Button("Example 3 - TravelAPI")
            
        with gr.Tab("Dual Model Demo"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("## Configurations")
                    dual_model_dropdown_1 = gr.Dropdown(choices=MODELS, label="Select Model 1", value=DEFAULT_MODEL_1, interactive=True)
                    dual_temperature_slider_1 = gr.Slider(0, 1, value=DEFAULT_TEMPERATURE_1, label="Temperature 1", interactive=True)
                    dual_model_dropdown_2 = gr.Dropdown(choices=MODELS, label="Select Model 2", value=DEFAULT_MODEL_2, interactive=True)
                    dual_temperature_slider_2 = gr.Slider(0, 1, value=DEFAULT_TEMPERATURE_2, label="Temperature 2", interactive=True)
                    dual_category_dropdown = gr.Dropdown(choices=CATEGORIES, label="Select Category", value=CATEGORIES[0], interactive=True)
                
                with gr.Column(scale=3):
                    with gr.Row():
                        with gr.Column():
                            gr.Markdown("### Model 1")
                            dual_chatbot1 = gr.Chatbot(elem_id="chatbot1", bubble_full_width=False, type="messages")
                            with gr.Row():
                                dual_restart_btn_1 = gr.Button("Restart")
                                dual_report_btn_1 = gr.Button("Report Issue")
                        with gr.Column():
                            gr.Markdown("### Model 2")
                            dual_chatbot2 = gr.Chatbot(elem_id="chatbot2", bubble_full_width=False, type="messages")
                            with gr.Row():
                                dual_restart_btn_2 = gr.Button("Restart")
                                dual_report_btn_2 = gr.Button("Report Issue")
                    with gr.Row():
                        dual_target_dropdown = gr.Dropdown(choices=["Model 1", "Model 2", "Both"], container=False, value="Both", interactive=True, scale=1)      
                        dual_chat_input = gr.MultimodalTextbox(
                            interactive=True,
                            file_count="multiple",
                            placeholder="Enter message or upload file...",
                            show_label=False,
                            scale=4
                        )

                    demo.load(lambda: get_initial_state(), outputs=dual_chatbot1)
                    demo.load(lambda: get_initial_state(), outputs=dual_chatbot2)
                    
                    dual_chat_msg = dual_chat_input.submit(
                        state_manager.add_message, [dual_chat_input, dual_target_dropdown], [dual_chatbot1, dual_chatbot2, dual_chat_input]
                    ) 

                    dual_bot_msg = dual_chat_msg.then(state_manager.get_reponse_dual, inputs=[dual_target_dropdown], outputs=[dual_chatbot1, dual_chatbot2])
                    
                    dual_bot_msg.then(lambda: gr.MultimodalTextbox(interactive=True), None, [dual_chat_input])

                    dual_chatbot1.like(print_like_dislike, None, None, like_user_message=True)
                    dual_chatbot2.like(print_like_dislike, None, None, like_user_message=True)
                    
                    with gr.Row():
                        dual_example_btn1 = gr.Button("Example 1 - GFSFileSystem")
                        dual_example_btn2 = gr.Button("Example 2 - TradingBot")
                        dual_example_btn3 = gr.Button("Example 3 - TravelAPI")
        
        with gr.Tab("Function Descriptions"):
            for category in CATEGORIES:
                with gr.Tab(category):
                    with gr.Group():
                        category_info = api_info[api_info["Class Name"] == category]
                        with gr.Accordion("Function description", open=False):
                            with gr.Group():
                                for i in range(len(category_info)):
                                    with gr.Accordion(category_info.iloc[i]["Function Name"], open=False):
                                        gr.Markdown(category_info.iloc[i]["Description"])

                        # Sample demo, limit 5 per categories
                        samples = [[sample['question'], sample['ground_truth']] for _, sample in api_samples.iterrows() if category in sample['involved_classes']][:5]
                        gr.Dataset(
                                components=[gr.HTML(), gr.Markdown()],
                                headers= ["Prompt", "API Use"],
                                samples= samples
                            )
                        
    single_model_dropdown.change(
        state_manager.single_bot.update_handler, 
        [single_model_dropdown, single_temperature_slider], 
        [single_model_dropdown, single_chatbot]
    )

    # Update handler when the temperature is changed
    single_temperature_slider.change(
        state_manager.single_bot.update_handler,
        [single_model_dropdown, single_temperature_slider], 
        [single_model_dropdown, single_chatbot]
    )

    # Update category and load config when a category is selected
    single_category_dropdown.change(
        state_manager.single_update_category_and_load_config,
        inputs=single_category_dropdown,
        outputs=single_category_dropdown
    )

    # Set up the event handler for the restart button to reset the chat and test_entry
    # restart_btn.click(restart_chat, None, [chatbot])
    single_restart_btn.click(state_manager.single_bot.restart_chat_and_save, [], [single_chatbot])
    single_report_btn.click(report_issue, None, None)
    
    single_example_btn1.click(state_manager.single_load_example_and_update, inputs=[single_example_btn1], 
                       outputs=[single_model_dropdown, single_temperature_slider, single_category_dropdown, single_chat_input])
    single_example_btn2.click(state_manager.single_load_example_and_update, inputs=[single_example_btn2], 
                       outputs=[single_model_dropdown, single_temperature_slider, single_category_dropdown, single_chat_input])
    single_example_btn3.click(state_manager.single_load_example_and_update, inputs=[single_example_btn3], 
                       outputs=[single_model_dropdown, single_temperature_slider, single_category_dropdown, single_chat_input])

    # Update handler when the model or temperature is changed
    dual_model_dropdown_1.change(
        state_manager.dual_bot1.update_handler, 
        [dual_model_dropdown_1, dual_temperature_slider_1], 
        [dual_model_dropdown_1, dual_chatbot1]
    )
    
    # Update handler when the model or temperature is changed
    dual_model_dropdown_2.change(
        state_manager.dual_bot2.update_handler, 
        [dual_model_dropdown_2, dual_temperature_slider_2], 
        [dual_model_dropdown_2, dual_chatbot2]
    )

    dual_temperature_slider_1.change(
        state_manager.dual_bot1.update_handler, 
        [dual_model_dropdown_1, dual_temperature_slider_1], 
        [dual_model_dropdown_1, dual_chatbot1]
    )

    dual_temperature_slider_2.change(
        state_manager.dual_bot2.update_handler, 
        [dual_model_dropdown_2, dual_temperature_slider_2], 
        [dual_model_dropdown_2, dual_chatbot2]
    )

    # Update category and load config when a category is selected
    dual_category_dropdown.change(
        state_manager.dual_update_category_and_load_config,
        inputs=dual_category_dropdown,
        outputs=dual_category_dropdown
    )

    # Set up the event handler for the restart button to reset the chat and test_entry
    dual_restart_btn_1.click(state_manager.dual_bot1.restart_chat_and_save, [], [dual_chatbot1])
    dual_report_btn_1.click(report_issue, None, None)
    
    dual_restart_btn_2.click(state_manager.dual_bot2.restart_chat_and_save, [], [dual_chatbot2])
    dual_report_btn_2.click(report_issue, None, None)
    
    dual_example_btn1.click(state_manager.dual_load_example_and_update, inputs=[dual_example_btn1], 
                       outputs=[dual_model_dropdown_1, dual_temperature_slider_1, dual_model_dropdown_2, dual_temperature_slider_2, dual_category_dropdown, dual_chat_input])
    dual_example_btn2.click(state_manager.dual_load_example_and_update, inputs=[dual_example_btn2], 
                       outputs=[dual_model_dropdown_1, dual_temperature_slider_1, dual_model_dropdown_2, dual_temperature_slider_2, dual_category_dropdown, dual_chat_input])
    dual_example_btn3.click(state_manager.dual_load_example_and_update, inputs=[dual_example_btn3], 
                       outputs=[dual_model_dropdown_1, dual_temperature_slider_1, dual_model_dropdown_2, dual_temperature_slider_2, dual_category_dropdown, dual_chat_input])

demo.launch(share=False)