pages/3_Customize_Holiday_Borders.py

import streamlit as st
import os
from PIL import Image
from io import BytesIO
import zipfile
from dotenv import load_dotenv

# Import from helper_utilities.py
from utils.helper_utilities import (
    get_next_largest_aspect_ratio, draw_crop_preview, 
    add_custom_message, generate_flux_image
)

# Import from configuration.py
from utils.configuration import (
    default_guidance_scale, default_num_inference_steps, 
    default_seed, holiday_border_prompts, example_holiday_messages, fonts
)

# Streamlit app starts here
st.image("img/fireworksai_logo.png")
st.title("🎄 Holiday Card Customizer 🎅")

st.markdown("""
✨ Welcome to the **Holiday Card Generator**! Get ready to sprinkle some holiday magic onto your custom card with a festive border and a heartfelt message. Here's how you can spread the holiday cheer: 

1. 🎁 **Upload your image** – Select a photo or design to be the star of your holiday card.
2. ✨ **Design your holiday card border** – Choose from seasonal prompts or create your own unique design.
3. 💌 **Add your personal holiday message** – Make it meaningful, funny, or warm-hearted!
4. 📦 **Download your finished card** – Ready to send to friends and family.

Let’s bring your holiday greetings to life! 🎉
""")

# Initialize session state variables
if 'uploaded_file' not in st.session_state:
    st.session_state.uploaded_file = None
if 'generated_image' not in st.session_state:
    st.session_state.generated_image = None
if 'metadata' not in st.session_state:
    st.session_state.metadata = {}

st.divider()
st.subheader("🎅 Step 1: Load Your Fireworks API Key 🎄")

# Load API Key
# Define and ensure the .env directory and file exist
dotenv_path = os.path.join(os.path.dirname(__file__), '..', 'env', '.env')
os.makedirs(os.path.dirname(dotenv_path), exist_ok=True)

# Create the .env file if it doesn't exist
if not os.path.exists(dotenv_path):
    with open(dotenv_path, "w") as f:
        st.success(f"Created {dotenv_path}")

# Load environment variables from the .env file
load_dotenv(dotenv_path, override=True)

# Check if the Fireworks API key is set or blank
fireworks_api_key = os.getenv("FIREWORKS_API_KEY")

# Show the entire app but disable running parts if no API key
if not fireworks_api_key or fireworks_api_key.strip() == "":
    fireworks_api_key = st.text_input("Enter Fireworks API Key", type="password")

    # Optionally, allow the user to save the API key to the .env file
    if fireworks_api_key and st.checkbox("Save API key for future use"):
        with open(dotenv_path, "a") as f:
            f.write(f"FIREWORKS_API_KEY={fireworks_api_key}\n")
        st.success("API key saved to .env file.")
else:
    st.success(f"API key loaded successfully: partial preview {fireworks_api_key[:5]}")

# Step 1: Upload Image
st.subheader("🎁 Step 2: Upload Your Festive Picture 🎨")
st.markdown("""
It's time to pick your image! This will be the centerpiece of your holiday card.  
- **Pro Tip**: Choose something that brings out the warmth and joy of the season! 🎄🎅  
Upload a **PNG** or **JPEG** image to get started.
""")
uploaded_file = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
if uploaded_file is not None:
    st.session_state.uploaded_file = uploaded_file

# Process uploaded image
if st.session_state.uploaded_file is not None:
    try:
        original_image = Image.open(st.session_state.uploaded_file)
        st.image(original_image, caption="Uploaded Image", use_column_width=True)

        # Step 2: Crop Image
        st.subheader("✂️ Step 3: Trim Your Image Just Right! 🎄")
        st.markdown("""
        Adjust the crop sliders to capture the most festive part of your image.  
        A preview will show you how your holiday card is shaping up! ✨
        """)
        img_width, img_height = original_image.size
        col1, col2 = st.columns(2)
        with col1:
            x_pos = st.slider("X position (Left-Right)", 0, img_width, img_width // 4)
            crop_width = st.slider("Width", 10, img_width - x_pos, min(img_width // 2, img_width - x_pos))
        with col2:
            y_pos = st.slider("Y position (Up-Down)", 0, img_height, img_height // 4)
            crop_height = st.slider("Height", 10, img_height - y_pos, min(img_height // 2, img_height - y_pos))

        preview_image = draw_crop_preview(original_image.copy(), x_pos, y_pos, crop_width, crop_height)
        st.image(preview_image, caption="Crop Preview", use_column_width=True)

        # Step 3: Set Border Parameters and Generate Image
        st.subheader("🎨 Step 4: Design Your Holiday Border! 🎅")
        st.markdown("""
        Choose from a range of holiday themes like snowflakes, cozy fireplaces, or fireworks!  
        Want to create your own magical scene? Enter a **custom prompt** to make it truly unique! 🎄
        """)
        selected_prompt = st.selectbox("Choose a holiday-themed prompt or enter your own", options=["Custom"] + holiday_border_prompts)
        custom_prompt = st.text_input("Enter your custom prompt") if selected_prompt == "Custom" else ""
        prompt = custom_prompt if selected_prompt == "Custom" else selected_prompt

        guidance_scale = st.slider("Guidance Scale", min_value=0.0, max_value=20.0, value=default_guidance_scale, step=0.1)
        num_inference_steps = st.slider("Number of Inference Steps", min_value=1, max_value=100, value=default_num_inference_steps, step=1)
        seed = st.slider("Random Seed", min_value=0, max_value=1000, value=default_seed)

        if st.button("Generate Holiday Card"):
            if not prompt.strip():
                st.error("Please enter a prompt.")
            else:
                with st.spinner("Adding holiday magic..."):
                    aspect_ratio = get_next_largest_aspect_ratio(img_width, img_height)
                    generated_image = generate_flux_image(
                        model_path="flux-1-schnell-fp8",
                        api_key=fireworks_api_key,
                        prompt=prompt,
                        steps=num_inference_steps,
                        guidance_scale=guidance_scale,
                        seed=seed,
                        aspect_ratio=f"{aspect_ratio[0]}:{aspect_ratio[1]}"
                    )
                    
                    if generated_image is not None:
                        generated_image = generated_image.resize(original_image.size)
                        cropped_original = original_image.crop((x_pos, y_pos, x_pos + crop_width, y_pos + crop_height))
                        center_x = (generated_image.width - crop_width) // 2
                        center_y = (generated_image.height - crop_height) // 2
                        final_image = generated_image.copy()
                        final_image.paste(cropped_original, (center_x, center_y))
                        st.session_state.generated_image = final_image

                        # Save metadata
                        st.session_state.metadata = {
                            "Prompt": prompt,
                            "Guidance Scale": guidance_scale,
                            "Inference Steps": num_inference_steps,
                            "Seed": seed
                        }

                        st.image(final_image, caption="🎁 Final Holiday Card 🎁", use_column_width=True)

                        # Print the metadata below the image
                        st.markdown("### ✨ Card Metadata")
                        for key, value in st.session_state.metadata.items():
                            st.write(f"**{key}:** {value}")

        # Step 4: Add Custom Message
        if 'generated_image' in st.session_state and st.session_state.generated_image is not None:
            st.subheader("💌 Step 5: Add Your Heartfelt Holiday Message! 🎅")
            st.markdown("""
            Select a pre-written festive message or craft your own!  
            Customize the font, size, and color to make it pop! 🌟
            """)

            selected_message = st.selectbox("Choose an example holiday message or write your own", options=["Custom"] + example_holiday_messages)
            custom_message = st.text_input("Enter your holiday message") if selected_message == "Custom" else selected_message
            
            selected_font = st.selectbox("Choose a font for your message", options=list(fonts.keys()))
            font_size = st.slider("Adjust font size", min_value=1, max_value=500, value=100, step=5)
            # Slider to choose the max characters per line for text wrapping
            max_chars = st.slider("Max characters per line:", min_value=20, max_value=100, value=40)
            font_path = fonts[selected_font]
            bg_color = st.color_picker("Pick a background color for the text", "#FFFFFF")
            font_color = st.color_picker("Pick a font color", "#000000")
            alpha = st.slider("Background transparency (0: fully transparent, 255: fully opaque)", 0, 255, 200)
            position_vertical = st.radio("Vertical position", ["Top", "Center", "Bottom"])
            position_horizontal = st.radio("Horizontal position", ["Left", "Center", "Right"])

            

            if st.button("Generate Holiday Card with Message"):
                final_image_with_message = st.session_state.generated_image.copy()

                final_image_with_message = add_custom_message(
                    final_image_with_message,
                    custom_message,
                    font_path,
                    font_size,
                    position_vertical,
                    position_horizontal,
                    max_chars=max_chars,
                    bg_color=bg_color,
                    font_color=font_color,
                    alpha=alpha
                )

                st.image(final_image_with_message, caption="🎄 Holiday Card with Custom Message 🎄", use_column_width=True)

                # Save image and metadata to zip
                img_byte_arr = BytesIO()
                final_image_with_message.save(img_byte_arr, format="PNG")
                img_byte_arr.seek(0)

                # Create a zip file containing the image and metadata
                zip_buffer = BytesIO()
                with zipfile.ZipFile(zip_buffer, "w") as zf:
                    # Add the image to the zip
                    zf.writestr("holiday_card_with_message.png", img_byte_arr.getvalue())

                    # Add the metadata as a text file to the zip
                    metadata_str = "\n".join([f"{key}: {value}" for key, value in st.session_state.metadata.items()])
                    zf.writestr("metadata.txt", metadata_str)

                zip_buffer.seek(0)

                # Step 5: Download as zip
                st.download_button(
                    label="🎁 Download Holiday Card and Metadata as ZIP 🎁",
                    data=zip_buffer,
                    file_name="holiday_card_with_message.zip",
                    mime="application/zip"
                )

    except OSError:
        st.error("The uploaded image seems to be truncated or corrupted. Please try with a different image.")
else:
    st.warning("Please upload an image to get started.")


    # Footer Section
st.divider()
st.markdown(
    """
    Thank you for using the Holiday Card Generator powered by **Fireworks**! 🎉  
    Share your creations with the world and spread the holiday cheer!  
    Happy Holidays from the **Fireworks Team**. 💥  
    """
)