import PIL.Image
import gradio as gr
import base64
import time
import os
import google.generativeai as genai
from gradio import ClearButton

# Set Google API key 
genai.configure(api_key = os.environ['GOOGLE_API_KEY'])

# Create the Model
txt_model = genai.GenerativeModel('gemini-pro')
vis_model = genai.GenerativeModel('gemini-pro-vision')

# Image to Base 64 Converter
def image_to_base64(image_path):
    with open(image_path, 'rb') as img:
        encoded_string = base64.b64encode(img.read())
    return encoded_string.decode('utf-8')

# Function that takes User Inputs and displays it on ChatUI
def query_message(history,txt,img):
    if not img:
        history += [(txt,None)]
        return history
    base64 = image_to_base64(img)
    data_url = f"data:image/jpeg;base64,{base64}"
    history += [(f"{txt} ![]({data_url})", None)]
    return history

# Function that takes User Inputs, generates Response and displays on Chat UI
def llm_response(history,text,img):
    if not img:
        response = txt_model.generate_content(text)
        history += [(None,response.text)]
        return history

    else:
        img = PIL.Image.open(img)
        response = vis_model.generate_content([text,img])
        history += [(None,response.text)]
        return history

# Interface Code
with gr.Blocks() as app:
    # Markdown block for title and description
    gr.Markdown("# HOW TO USE")
    gr.Markdown("1. Take a pic of the exam or textbook diagram and upload here a copy of it from your mobile, system or device.\n\n2. Enter text describing what you want to learn about the image, e.g, 'Please carefully analyze the image I have uploaded and answer this question for me, step-by-step, with explanations/workings...'.\n\n3. Click 'Submit' and await a response from the AI or 'Clear' to upload a new image for analysis. ")
    gr.Markdown("""
    **Note**: While AI teaches you how to solve problems, it is not infallible or without potential to make mistakes or errors. So feel free to double-check answers with your textbook knowledge or schoolteacher's guidance.
    """)
    with gr.Row():
        image_box = gr.Image(type="filepath")
    
        chatbot = gr.Chatbot(
            scale = 2,
            height=750
        )
    text_box = gr.Textbox(
            placeholder="Tell me the question(s) you'd like me to answer about the image you just uploaded then click 'Submit' to get my answer. Click 'Clear' to upload and process a new image.",
            container=False,
        )

    btn = gr.Button("Submit")
    
    # Create a ClearButton instance with the components to clear
    clear_btn = ClearButton([chatbot, text_box, image_box])
    
    clicked = btn.click(query_message,
                        [chatbot,text_box,image_box],
                        chatbot
                        ).then(llm_response,
                                [chatbot,text_box,image_box],
                                chatbot
                                )

app.queue()
app.launch(debug=True)