import gradio as gr
from src.main import ChatWrapper

chat = ChatWrapper()

def update_chain(api_key: str, selection: str):
    global chat_wrapper
    chat_wrapper = ChatWrapper(chain=selection, api_key=api_key)

block = gr.Blocks(css=".gradio-container {background-color: lightgray}")

with block:
    with gr.Row():
        gr.Markdown("<h3><center>Hello-World LangChain App</center></h3>")
        selection = gr.Dropdown(label="Select Agent", choices=["falcon", "openai"], default="openai")
        api_key_textbox = gr.Textbox(
            label="API Key",
            placeholder="Paste your OpenAI API key (sk-...)",
            show_label=True,
            lines=1,
            type="password",
        )

    chatbot = gr.Chatbot()

    with gr.Row():
        message = gr.Textbox(
            label="What's your question?",
            placeholder="What's the answer to life, the universe, and everything?",
            lines=1,
        )
        submit = gr.Button(value="Send", variant="secondary").style(full_width=False)

    gr.Examples(
        examples=[
            "Hi! How's it going?",
            "What should I do tonight?",
            "Whats 2 + 2?",
        ],
        inputs=message,
    )

    gr.HTML("Demo application of a LangChain chain.")

    state = gr.State()
    agent_state = gr.State()

    submit.click(chat, inputs=[api_key_textbox, message, state, agent_state], outputs=[chatbot, state])
    message.submit(chat, inputs=[api_key_textbox, message, state, agent_state], outputs=[chatbot, state])

    api_key_textbox.change(update_chain, inputs=[api_key_textbox, selection], outputs=[agent_state])

block.launch(debug=True)