import gradio as gr
from datetime import datetime

from chatbot import Chatbot
from client_utils import engine_map

chatbot_engine = Chatbot()

def get_user(request: gr.Request):
    try:
        return request.headers.user
    except:
        return ""

def generate_greeting(user) -> str:
    hour = datetime.now().hour
    if 5 <= hour < 12:
        greeting = "Good morning"
    elif 12 <= hour < 17:
        greeting = "Good afternoon"
    else:
        greeting = "Good evening"
    return f'<p style="font-size: 24px; text-align: center;"><span style="font-weight: bold;">{greeting}{f", {user}" if user else ""}.</span><br>How can I help you?</p>'

theme = gr.themes.Monochrome(
    radius_size="xxl",
    font=['Montserrat', 'ui-sans-serif', 'system-ui', 'sans-serif'],
).set(
    background_fill_primary='white',
    background_fill_secondary='white',
    background_fill_secondary_dark='*neutral_950',
    block_background_fill='*neutral-50',
    block_border_width='0px',
    block_border_width_dark='0px',
    block_label_border_width='0px',
    block_label_border_width_dark='0px',
    border_color_accent='white',
    border_color_accent_dark='*neutral_950',
)

with gr.Blocks(theme=theme, css_paths="style.css", fill_height=True) as demo:
    user = gr.State()
    
    with gr.Sidebar():
        chat_history_dataset = gr.Dataset(
            components=[gr.Textbox(visible=False)],
            samples=[[]],
            label="Recent",
            show_label=True,
            layout="table",
            type="index",
        )
    
    with gr.Row():
        with gr.Column():
            engine = gr.Dropdown(  # Renamed from "model" to "engine"
                choices=list(chatbot_engine.engine_map.keys()),  # Updated to use engine_map
                value=list(chatbot_engine.engine_map.keys())[-1],
                show_label=False,
            )
        with gr.Column():
            new_chat_btn = gr.Button(
                "",
                icon="assets/new_chat.svg",
                variant="primary",
                elem_id="new_chat_btn"
            )
            login_btn = gr.LoginButton(
                "Login",
                logout_value="Logout",
                icon="assets/login.svg",
                variant="primary",
                elem_id="login_btn"
            )
        
    chatbot = gr.Chatbot(
        type="messages",
        height="100%",
        max_height= "75vh",
        show_label=False,
        elem_id="chatbot",
        placeholder=generate_greeting("")
    )
    
    with gr.Group(elem_id="inputs", render=False) as inputs:
        textbox = gr.MultimodalTextbox(
            placeholder="Message Digichat...",
            file_count="multiple",
            show_label=False,
        )
        with gr.Row(visible=False):
            temperature = gr.Slider(
                minimum=0,
                maximum=1,
                step=0.1,
                value=0.0,
                show_label=False,
            )
        engine_params = [engine, temperature]
    
    chat_interface = gr.ChatInterface(
        chatbot_engine.predict,
        type="messages",
        chatbot=chatbot,
        textbox=textbox,
        additional_inputs=engine_params, 
        additional_outputs=[textbox],
        editable=True,
        save_history=True,
    )
    
    inputs.render()
    
    synchronize_chat_state_kwargs = {
        "fn": lambda x: (x, x),
        "inputs": [chat_interface.chatbot],
        "outputs": [chat_interface.chatbot_state, chat_interface.chatbot_value],
        "show_api": False,
        "queue": False,
    }

    new_chat_btn.click(
        lambda: (None, []),
        None,
        [chat_interface.conversation_id, chat_interface.chatbot],
        show_api=False,
        queue=False,
    ).then(
        lambda x: x,
        [chat_interface.chatbot],
        [chat_interface.chatbot_state],
        show_api=False,
        queue=False,
    )

    gr.on(
        triggers=[chat_interface.load, chat_interface.saved_conversations.change],
        fn=chat_interface._load_chat_history,
        inputs=[chat_interface.saved_conversations],
        outputs=[chat_history_dataset],
        show_api=False,
        queue=False,
    )

    chat_history_dataset.click(
        lambda: [],
        None,
        [chat_interface.chatbot],
        show_api=False,
        queue=False,
        show_progress="hidden",
    ).then(
        chat_interface._load_conversation,
        [chat_history_dataset, chat_interface.saved_conversations],
        [chat_interface.conversation_id, chat_interface.chatbot],
        show_api=False,
        queue=False,
        show_progress="hidden",
    ).then(**synchronize_chat_state_kwargs)

    gr.on(
        [textbox.stop, chatbot.clear, new_chat_btn.click],
        lambda: chatbot_engine.stop(),
        show_progress='hidden'
    )

    demo.load(
        get_user,
        None,
        [user]  
    ).then(
        lambda user: gr.update(placeholder=generate_greeting(user)),
        [user],
        [chatbot]
    )

if __name__ == "__main__":
    demo.queue(
        default_concurrency_limit=40
    ).launch(
        pwa=True, 
        share=False
    )