import streamlit as st
from PIL import Image, ImageOps
from internal.api import APIClient
import os

client = APIClient("https://collage-ai.onrender.com")
# client = APIClient("http://localhost:3000")

def gallery(column, images):
    groups = []
    for i in range(0, len(images), column):
        groups.append(images[i:i+column])

    for group in groups:
        cols = st.columns(column)
        for i, image in enumerate(group):
            cols[i].image(image)

def resize_image(image_file):
    image = Image.open(image_file)
    resized = image.resize((224, 224))
    print(f"Resized image: {resized}")
    resized.save(image_file.name)
    return open(image_file.name, 'rb')

st.title('CollageAI')

user_images = None
user_prompt = None
uploaded_images = None
captions = []
keywords = []
submitted = False

with st.form("user_input_form"):
    user_images = st.file_uploader(
        "Choose your photos",
        accept_multiple_files=True
    )
    user_prompt = st.text_area(
        "Describe the design you'd like to create:",
        placeholder="For our anniversary, I want to write a card to my partner to celebrate our love and share all the things I adore about them."
    )
    submitted = st.form_submit_button("Generate")

# Check form
if submitted:
    if user_images:
        with st.container():
            with st.spinner('Uploading images...'):
                try:
                    resized_images = [resize_image(image) for image in user_images]
                    uploaded_images = client.upload_images(resized_images)
                    # delete resized image files
                    for image in resized_images:
                        image.close()
                        os.remove(image.name)
                except Exception as e:
                    st.error(f"Error uploading images: {e}")

            # Display the photo gallery
            st.subheader('Your photos:')
            images = [Image.open(image) for image in user_images]
            images = [ImageOps.exif_transpose(image) for image in images]
            gallery(4, images)
    else:
        st.warning('Please upload at least one image before submitting.')

    if user_prompt:
        if uploaded_images:
            # Analysis
            with st.spinner('Analyzing prompt...'):
                try:
                    analysis = client.analyze_prompt(user_prompt, uploaded_images)
                    keywords = analysis.get("keywords")
                    captions = analysis.get("captions")
                    if captions:
                        st.subheader('Captions of your photos')
                        st.write(captions)

                    if keywords:
                        st.subheader('Keywords based on your photos and prompt')
                        st.write(keywords)
                    else:
                        st.warning('No keywords were generated. Please try again with a different prompt.')

                except Exception as e:
                    st.error(f"Error analyzing prompt: {e}")

        # Stickers
        with st.container():
            with st.spinner('Generating stickers...'):
                try:
                    sticker_image_urls = client.suggest_stickers(user_prompt, captions)
                    if sticker_image_urls:
                        st.subheader('Stickers suggestions')
                        gallery(4, sticker_image_urls[:8])
                    else:
                        st.warning('No images were generated. Please try again with a different prompt.')
                except Exception as e:
                    st.error(f"Error generating stickers: {e}")

        # Templates
        with st.container():
            with st.spinner('Generating templates...'):
                try:
                    template_image_urls = client.suggest_templates(user_prompt, captions)
                    if template_image_urls:
                        st.subheader('Template suggestions')
                        gallery(4, template_image_urls[:8])
                    else:
                        st.warning('No images were generated. Please try again with a different prompt.')
                except Exception as e:
                    st.error(f"Error generating templates: {e}")

    else:
        st.warning('Please enter a prompt before submitting.')