# Generate_Holiday_Postcard.py

import streamlit as st
from utils.helper_utilities import generate_flux_image, add_custom_message
from utils.configuration import fonts, holiday_scene_prompts, example_holiday_messages
from dotenv import load_dotenv
import os

# Set page configuration
st.set_page_config(page_title="FLUX Image Generation Tool", page_icon="🎇")

# Streamlit UI Elements
st.title("🎇 FLUX-tastic Holiday Postcard Generator 🎨")
st.markdown(
    """Welcome to the FLUX Holiday Postcard Generator! 

- 🎅🎨 It's time to make your holiday greetings pop with a personalized, AI-generated postcard! 
- ✨ No more boring, store-bought cards! Instead, flex your creative muscles with the FLUX models, adding your unique touch. 
- Customize a dazzling holiday scene with your own message and share your festive vibes! 🎁🎄

### How it works:

1. Choose a holiday-themed prompt or write your own 🎉
2. Select a FLUX model to bring your vision to life ✨
3. Customize with a snazzy holiday message and choose from a variety of fonts 🖋️ 
4. Design your text's background and color to perfectly match the vibe 🎨 
5. Generate your masterpiece and share the joy with friends and family! 📬

Get ready to deck the halls with creativity! 🎄✨

"""
)

st.divider()
st.subheader("Load 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]}")


# Dropdown to select a holiday-themed prompt or enter a custom prompt
st.divider()
st.subheader("1️⃣ Step 1: Pick Your Holiday Theme or Create Your Own 🎄✨")
st.markdown(""" 
Get into the festive spirit by choosing from a range of holiday-inspired prompts. Feeling extra creative? Enter your own prompt and let the holiday magic begin! 🎅✨
""")

selected_prompt = st.selectbox("Choose a holiday-themed prompt or enter your own", options=["Custom"] + holiday_scene_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

# Dropdown to select the model
st.divider()
st.subheader("2️⃣ Step 2: Select Your FLUX Model 🚀")
st.markdown(""" 
Choose from two FLUX models: whether you’re aiming for lightning speed or extra detail, we’ve got you covered! 💥✨
""")

model_choice = st.selectbox("Select the model:", ["flux-1-schnell-fp8", "flux-1-dev-fp8"])

# Expose parameters like guidance_scale, inference steps, and seed after model selection
st.divider()
st.subheader("2️⃣A: Customize Your Model Parameters ⚙️")
st.markdown("""
Now that you've selected your model, fine-tune the parameters to adjust how the FLUX model generates your holiday card. 
""")

# Parameter options exposed after model selection
guidance_scale = st.slider("Guidance Scale", min_value=0.0, max_value=20.0, value=7.5, step=0.1)
num_inference_steps = st.slider("Number of Inference Steps", min_value=1, max_value=100, value=50, step=1)
seed = st.slider("Random Seed", min_value=0, max_value=1000, value=42)


# Dropdown to select an example message or write a custom one
st.divider()
st.subheader("3️⃣ Step 3: Craft the Perfect Message 🎁")
st.markdown(""" 
What’s a holiday card without a heartfelt (or hilarious) message? Choose from example holiday greetings or write your own to make the card so you! 🖋️🎉
""")

selected_message = st.selectbox("Choose an example message or write your own:", options=["Custom"] + example_holiday_messages)
message = st.text_input("Enter a holiday message to add:", value=selected_message if selected_message != "Custom" else "")

# Additional inputs for customizing the message
st.divider()
st.subheader("4️⃣ Step 4: Style Your Message with Flair ✨")
st.markdown(""" 
From fancy fonts to colorful backgrounds, you’re in control. Pick your favorite font, adjust the size, and add a splash of color to make your message truly shine. 🌈🎨
""")

font_choice = st.selectbox("Select a font:", list(fonts.keys()))
font_size = st.slider("Select font size:", 1, 300, 40)
max_chars = st.slider("Max characters per line:", 10, 100, 40)  # Slider to select character wrap limit

# Background color and font color pickers
bg_color = st.color_picker("Pick a background color for the text box:", "#FFFFFF")
font_color = st.color_picker("Pick a font color:", "#000000")

# Transparency level slider
alpha = st.slider("Select transparency level for the background (0: fully transparent, 255: fully opaque)", 0, 255, 220)

# Position options for the message (vertical and horizontal)
position_vertical = st.radio("Select message vertical position:", ["Top", "Center", "Bottom"])
position_horizontal = st.radio("Select message horizontal position:", ["Left", "Center", "Right"])

st.divider()
st.subheader("5️⃣ Step 5: Preview and Share the Holiday Cheer! 🎅📬")
st.markdown(""" 
Click "Generate Image" and watch the magic happen! Your holiday card is just moments away from spreading joy to everyone on your list. 🎄🎁✨
""")

# Button to generate images
if not fireworks_api_key or fireworks_api_key.strip() == "":
    st.warning("Enter a valid Fireworks API key to enable image generation.")
    generate_button = st.button("Generate Image", disabled=True)
else:
    generate_button = st.button("Generate Image")


if generate_button:
    st.markdown(""" 
        🎉 You're one click away from holiday magic! 🎉 Hit that Generate Image button and let FLUX create your personalized postcard—ready for sharing! 📬
        """)
    if not prompt.strip():
        st.error("Please provide a prompt.")
    else:
        try:
            with st.spinner("Generating image..."):
                # Determine steps based on model
                steps = 30 if model_choice == "flux-1-dev-fp8" else 4

                # Generate image using the helper utility
                generated_image = generate_flux_image(
                    model_path=model_choice,        # Correctly pass the model choice
                    api_key=fireworks_api_key,      # Fireworks API key
                    prompt=prompt,                  # User's prompt for generation
                    steps=steps,                    # Number of inference steps
                    guidance_scale=guidance_scale,  # Guidance scale (optional if unchanged)
                    seed=seed               # Random seed (optional if unchanged)
)

                # Get the selected font path from configuration
                font_path = fonts[font_choice]

                # Add the holiday message to the generated image using the helper utility
                image_with_message = add_custom_message(
                    generated_image.copy(), message, font_path, font_size, 
                    position_vertical, position_horizontal, max_chars, bg_color, font_color, alpha
                )

                # Display the image with the message
                st.image(image_with_message, caption=f"Generated using {model_choice} with custom message", use_column_width=True)

                # Preview the placement details
                st.write(f"Message preview (vertical position: {position_vertical}, horizontal position: {position_horizontal}, font: {font_choice}, size: {font_size}, max chars: {max_chars}, bg color: {bg_color}, font color: {font_color}, transparency: {alpha})")
                
        except Exception as e:
            st.error(f"An error occurred: {e}")

# 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**. 💥  
    """
)