import streamlit as st
import os
from PIL import Image
import io
from utils.gemini_utils import get_meme_template_description
from utils.meme_utils import generate_meme_text

def main():
    st.title("Meme Generator")

    # Meme template gallery
    st.header("Meme Template Gallery")
    meme_templates = [f for f in os.listdir("meme_templates") if f.endswith(('.png', '.jpg', '.jpeg'))]
    
    # Use session state to store the selected template
    if 'selected_template' not in st.session_state:
        st.session_state.selected_template = None

    # Create a 4-column layout
    cols = st.columns(4)
    
    for idx, template in enumerate(meme_templates):
        col = cols[idx % 4]
        image_path = os.path.join("meme_templates", template)
        with Image.open(image_path) as img:
            col.image(img, use_column_width=True)
        if col.button(f"Select {template}", key=f"btn_{template}"):
            st.session_state.selected_template = template

    # Display selected template and generate meme
    if st.session_state.selected_template:
        st.subheader(f"Selected Template: {st.session_state.selected_template}")
        image_path = os.path.join("meme_templates", st.session_state.selected_template)
        with Image.open(image_path) as img:
            st.image(img, caption=st.session_state.selected_template, use_column_width=True)

        # Get user input for topic
        topic = st.text_input("Enter a topic for your meme:")

        if st.button("Generate Meme"):
            with st.spinner("Generating meme..."):
                # Get meme template description
                with open(image_path, "rb") as img_file:
                    img_bytes = img_file.read()
                
                meme_description = get_meme_template_description(img_bytes)

                # Generate meme text
                meme_text = generate_meme_text(meme_description, topic)

                # Display results
                st.subheader("Generated Meme Text:")
                st.write(meme_text)

                st.subheader("Meme Template Description:")
                st.write(meme_description)

if __name__ == "__main__":
    main()