import gradio as gr
from src.main import ChatWrapper

agent = ChatWrapper('openai', '')  # default agnet_state

def update_agent(api_key: str, selection: str):
    global agent
    agent = ChatWrapper(chain_type=selection, api_key=api_key)
    return agent # This is agent state

def chat(message):
    global agent
    agent(message)  # Get a response to the current message
    history = agent.history  # Access the entire chat history
    return history, history # Return the history twice to update both the chatbot and the state

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

with block:
    with gr.Row():
        gr.HTML("<h2><center>ConversationalChain App 🤖</center></h2> <br> For <a href='https://huggingface.co/tiiuae/falcon-7b-instruct'>Falcon</a> use HuggingFace API Token, for OpenAI use OpenAI API Key")
        selection = gr.Dropdown(label="Select Agent", choices=["falcon", "openai"], default="openai")        
        api_key_textbox = gr.Textbox(
            label="API Key",
            placeholder= "Paste Your API Key",
            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("<center>View more at  <a href='https://ai.rohankataria.com'>ai.rohankataria.com</a></center>")

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

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

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

block.launch(debug=True)