import streamlit as st
from PIL import Image
import textwrap
import google.generativeai as genai

# Function to display formatted Markdown text
def to_markdown(text):
    text = text.replace('•', '  *')
    return textwrap.indent(text, '> ', predicate=lambda _: True)

# Function to generate content using Gemini API
def generate_gemini_content(prompt, model_name='gemini-pro', image=None):
    model = genai.GenerativeModel(model_name)
    if image:
        response = model.generate_content([prompt, image])
    else:
        response = model.generate_content(prompt)

    return response

# Streamlit app
def main():
    st.title("Gemini API Demo with Streamlit")

    # Get Gemini API key from user input
    api_key = st.text_input("Enter your Gemini API key:")
    genai.configure(api_key=api_key)

    # Choose a model
    model_name = st.selectbox("Select a Gemini model", ["gemini-pro", "gemini-pro-vision"])

    # Get user input prompt
    prompt = st.text_area("Enter your prompt:")

    # Get optional image input
    image_file = st.file_uploader("Upload an image (if applicable):", type=["jpg", "jpeg", "png"])

    # Display image if provided
    if image_file:
        st.image(image_file, caption="Uploaded Image", use_column_width=True)

    # Generate content on button click
    if st.button("Generate Content"):
        st.markdown("### Generated Content:")
        if image_file:
            # If an image is provided, use gemini-pro-vision model
            image = Image.open(image_file)
            response = generate_gemini_content(prompt, model_name='gemini-pro-vision', image=image)
        else:
            response = generate_gemini_content(prompt, model_name=model_name)

        # Display the generated content in Markdown format
        st.markdown(to_markdown(response.text))

if __name__ == "__main__":
    main()