import streamlit as st
from streamlit_option_menu import option_menu
from utility import *
from PIL import Image

st.set_page_config(
    page_title="Gemini AI bot",
    page_icon="🔮",
    layout="centered",

)

with st.sidebar:
    selected = option_menu("Gemini AI bot", [
                           "Image Captioning", "Embed text", "Ask Anything"],
                           menu_icon="robot",
                           icons=["image-fill", "textarea-t",
                                  "patch-question-fill"],
                           default_index=0)


def map_model_role_to_streamlit_role(user_role):
    if user_role == "model":
        return "assistant"
    else:
        return user_role


model = load_model()

if selected == "Image Captioning":
    st.title("📸 Image Captioning")

    uploaded_image = st.file_uploader(
        "Upload an image", type=["jpg", "jpeg", "png"])

    if st.button("Generate caption"):
        image = Image.open(uploaded_image)
        col1, col2 = st.columns(2)

        with col1:
            resized_image = image.resize((1024, 720))
            st.image(resized_image, caption="Original image")

        default_prompt = "Write a short caption for this image"

        # Pass the image bytes to the model
        response = model.generate_content([image, default_prompt])
        caption = response.text

        with col2:
            st.info(caption)



if selected == "Embed text":
    st.title("🔡 Embed text")
    
    input_text = st.text_area(
        label="", placeholder="Enter text here to get the embeddings")
    
    if st.button("Get embeddings"):
        response = embedding_model_response(input_text)
        st.markdown(response)


if selected == "Ask Anything":
    st.title("🤔 Ask Anything")

    user_prompt = st.text_area(
        label="", placeholder="Enter your question here")

    if st.button("Get answer"):
        response = gemini_pro_response(user_prompt)
        st.write(response)