import os
import gradio as gr
import google.generativeai as genai


GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")

genai.configure(api_key=GEMINI_API_KEY)
model = genai.GenerativeModel('gemini-pro')

header = r"""
<center>
    <h1>Gemini Chat</h1>
    <h2>
        This demo is for code reference, feel free to clone it and set your own Api key
    </h2>
    Get an Api Key: <a href='https://aistudio.google.com/app/apikey'>Google Ai Studio</a>
</center>
"""

footer = r"""
<center>
    <b>
        Demo for <a href='https://gemini.google.com/app'>Google Gemini</a>
    </b>
</center>
"""

with gr.Blocks(title="Gemini Chat") as app:

    gr.HTML(header)

    with gr.Row(equal_height=True):
        with gr.Column():
            chatbot = gr.Chatbot()
            msg = gr.Textbox()
            clear = gr.Button("Clear")

            def user(user_message, history):
                return "", history + [[user_message, None]]

            def bot(history):
                response = model.generate_content(history[-1][0], stream=True)
                history[-1][1] = ""
                for chunk in response:
                    history[-1][1] += chunk.text
                    yield history

            msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
                bot, chatbot, chatbot
            )
            clear.click(lambda: None, None, chatbot, queue=False)
    
    with gr.Row():
        gr.HTML(footer)
    
app.queue()
app.launch(share=False, debug=True, show_error=True)