import streamlit as st
from google import genai
from google.genai import types
from PIL import Image
from io import BytesIO
import tempfile
import os
from dotenv import load_dotenv
load_dotenv()
GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY")
os.environ["GOOGLE_API_KEY"] = GOOGLE_API_KEY

client = genai.Client(api_key=GOOGLE_API_KEY)

st.set_page_config(page_title="Image Generator", layout="centered")
st.title("🧠✨ Image Generator")
st.markdown("Enter a prompt below and generate an AI image with a description!")

prompt = st.text_input("Enter your prompt")

if st.button("Generate"):
    if prompt:
        with st.spinner("Generating image and description..."):
            try:
                response = client.models.generate_content(
                    model="gemini-2.0-flash-exp-image-generation",
                    contents=prompt,
                    config=types.GenerateContentConfig(
                        response_modalities=['Text', 'Image']
                    )
                )

                result_text = ""
                result_image = None
                temp_file_path = None

                for part in response.candidates[0].content.parts:
                    if part.text is not None:
                        result_text += part.text
                    elif part.inline_data is not None:
                        result_image = Image.open(BytesIO(part.inline_data.data))
                        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
                        result_image.save(temp_file.name)
                        temp_file_path = temp_file.name

                st.subheader("Generated Description")
                st.text(result_text)

                if result_image:
                    st.image(result_image, caption="Generated Image", use_column_width=True)
                    with open(temp_file_path, "rb") as f:
                        st.download_button(
                            label="📥 Download Image",
                            data=f,
                            file_name="generated_image.png",
                            mime="image/png"
                        )

            except Exception as e:
                st.error(f"An error occurred: {e}")
    else:
        st.warning("Please enter a prompt before generating.")