import streamlit as st
import google.generativeai as genai
import os

# Securely fetch the API key from environment variables
api_key = os.getenv("GEMINI_API_KEY")

if not api_key:
    st.error("API key not found. Please set the GEMINI_API_KEY environment variable.")
    st.stop()

# Set up Gemini API
genai.configure(api_key=api_key)

# Set up the model
generation_config = {
    "temperature": 0.9,  # Controls creativity (0 = strict, 1 = creative)
    "top_p": 1,
    "top_k": 1,
    "max_output_tokens": 2048,  # Maximum length of the response
}

safety_settings = [
    {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
    {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
    {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
    {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
]

model = genai.GenerativeModel(
    model_name="gemini-2.0-pro-exp-02-05",
    generation_config=generation_config,
    safety_settings=safety_settings,
)

# Streamlit App
st.title("📖 Gemini Story Generator")
st.write("Enter a prompt, and Gemini AI will generate a story for you!")

# User input
user_prompt = st.text_area("Enter your story prompt:", "Once upon a time...")

# Generate story button
if st.button("Generate Story"):
    if user_prompt.strip() == "":
        st.warning("Please enter a prompt!")
    else:
        try:
            # Generate response using Gemini AI
            response = model.generate_content(
                f"Write a creative and engaging story suitable for elementary school children based on this prompt: {user_prompt}"
            )
            st.write("### Here's Your Story:")
            st.write(response.text)
            st.success("🌟 The End 🌟")
        except Exception as e:
            st.error(f"An error occurred: {e}. Please try again.")