import gradio as gr
from chatllm import (
    chat_response, get_llm_model, set_llm_model, get_llm_model_info,
    get_llm_language, set_llm_language, get_llm_sysprompt, get_llm_sysprompt_mode,
    set_llm_sysprompt_mode, 
)

# Custom CSS for Gradio app
css = '''
.gradio-container{max-width: 1000px !important}
h1{text-align:center}
footer { visibility: hidden }
'''

# Create Gradio interface
with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", css=css) as app:
    with gr.Column():
        with gr.Group():
            chatbot = gr.Chatbot(likeable=False, show_copy_button=True, show_share_button=False, layout="bubble", container=True)
            with gr.Row():
                chat_query = gr.Textbox(label="Search Query", placeholder="hatsune miku", value="", scale=3)
                chat_clear = gr.Button("🗑️ Clear", scale=1)
            with gr.Row():
                chat_msg = gr.Textbox(label="Message", placeholder="Input message with or without query and press Enter or click Sumbit.", value="", scale=3)
                chat_submit = gr.Button("Submit", scale=1)
            with gr.Accordion("Additional inputs", open=False):
                chat_model = gr.Dropdown(choices=get_llm_model(), value=get_llm_model()[0], allow_custom_value=True, label="Model")
                chat_model_info = gr.Markdown(value=get_llm_model_info(get_llm_model()[0]), label="Model info")
                with gr.Row():
                    chat_mode = gr.Dropdown(choices=get_llm_sysprompt_mode(), value=get_llm_sysprompt_mode()[0], allow_custom_value=False, label="Mode")
                    chat_lang = gr.Dropdown(choices=get_llm_language(), value="language same as user input", allow_custom_value=True, label="Output language")
                chat_tokens = gr.Slider(minimum=1, maximum=4096, value=2000, step=1, label="Max tokens")
                chat_temp = gr.Slider(minimum=0.1, maximum=4.0, value=0.9, step=0.1, label="Temperature")
                chat_topp = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p")
                chat_fp = gr.Slider(minimum=0.0, maximum=2.0, value=0.0, step=0.1, label="Frequency penalty")
                chat_sysmsg = gr.Textbox(value=get_llm_sysprompt(), interactive=True, label="System message")
        examples = gr.Examples(
            examples = [
                ["Describe this person.", "Kafuu Chino from Gochiusa"],
                ["Hello", ""],
            ],
            inputs=[chat_msg, chat_query],
        )
        gr.Markdown(
            f"""This demo was created in reference to the following demos.<br>
    [prithivMLmods/WEB-DAC](https://huggingface.co/spaces/prithivMLmods/WEB-DAC), 
    [featherless-ai/try-this-model](https://huggingface.co/spaces/featherless-ai/try-this-model), 
                """
        )
        gr.DuplicateButton(value="Duplicate Space")
        gr.Markdown(f"Just a few edits to *model.py* are all it takes to complete your own collection.")
        gr.on(
            triggers=[chat_msg.submit, chat_query.submit, chat_submit.click],
            fn=chat_response,
            inputs=[chat_msg, chatbot, chat_query, chat_tokens, chat_temp, chat_topp, chat_fp],
            outputs=[chatbot],
            queue=True,
            show_progress="full",
            trigger_mode="once",
        )
        chat_clear.click(lambda: (None, None, None), None, [chatbot, chat_msg, chat_query], queue=False)
        chat_model.change(set_llm_model, [chat_model], [chat_model, chat_model_info], queue=True, show_progress="full")\
        .success(lambda: None, None, chatbot, queue=False)
        chat_mode.change(set_llm_sysprompt_mode, [chat_mode], [chat_sysmsg], queue=False)
        chat_lang.change(set_llm_language, [chat_lang], [chat_sysmsg], queue=False)

if __name__ == "__main__":
    app.queue()
    app.launch()