import os
import streamlit as st
from PIL import Image
import google.generativeai as genai
from constants import gemini_key

# Streamlit framework configuration
st.set_page_config(
    page_title="OxSecure Images",
    page_icon="🎨",
    layout="wide"
)

# API configuration
os.environ["GOOGLE_API_KEY"] = gemini_key
genai.configure(api_key=os.environ['GOOGLE_API_KEY'])

# Function to load Gemini vision model and get responses
def get_gemini_response(input_text, image=None):
    model = genai.GenerativeModel('gemini-pro-vision')
    if image is not None:
        response = model.generate_content([input_text, image])
    else:
        response = model.generate_content(input_text)
    return response.text

def generate_gemini_image(prompt, num_images=1, size="1024x1024"):
    model = genai.GenerativeModel('gemini-pro-vision')
    # Placeholder for image generation method; replace with actual method
    response = model.generate_content(prompt)  # This line should be replaced with the correct method for image generation
    # Assuming the API returns a list of images as URLs or base64 encoded strings
    return response.images[:num_images]  # Adjust this line based on actual API response format

# Streamlit Main Framework
st.header('OxSecure ImaGen 🎨')
st.title('GenAI ImaGen powers ♨️')
st.subheader('By :- Aadi 🧑‍💻')

# Text input for prompt
input_text = st.text_input("Input Prompt: ", key="input")

# File uploader for image
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
image = None
if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image.", use_column_width=True)

# Button to get response about the image
submit_analyze = st.button("Tell me about the image")
if submit_analyze:
    if input_text and image is not None:
        response = get_gemini_response(input_text, image)
    elif image is not None:
        response = get_gemini_response("", image)
    elif input_text:
        response = get_gemini_response(input_text)
    else:
        response = "Please provide an input prompt or upload an image."
    st.subheader("The Response is")
    st.write(response)

# Button to generate an image from a prompt
submit_generate = st.button("Generate Image from Prompt")
if submit_generate and input_text:
    num_images = st.number_input("Number of Images to Generate", min_value=1, max_value=5, value=1)
    size = st.selectbox("Select Image Size", ["512x512", "1024x1024", "2048x2048"], index=1)
    generated_images = generate_gemini_image(input_text, num_images=num_images, size=size)
    for img in generated_images:
        st.image(img, caption="Generated Image", use_column_width=True)
else:
    if not input_text:
        st.write("Please provide an input prompt to generate an image.")