# app.py
import streamlit as st
from core.gemini_handler import GeminiHandler
from core.visual_engine import VisualEngine
from core.prompt_engineering import create_story_breakdown_prompt, create_image_prompt_from_scene
import os

# --- Configuration & Initialization ---
st.set_page_config(page_title="CineGen AI", layout="wide")

# Load API Key from Streamlit secrets
try:
    GEMINI_API_KEY = st.secrets["GEMINI_API_KEY"]
except KeyError:
    st.error("GEMINI_API_KEY not found in secrets. Please add it to your Streamlit Cloud secrets.")
    st.stop()

gemini = GeminiHandler(api_key=GEMINI_API_KEY)
visualizer = VisualEngine(output_dir="temp_generated_media") # Use a temp dir for HF Spaces

# --- UI Elements ---
st.title("🎬 CineGen AI: Your Pocket AI Film Studio")
st.markdown("Let's turn your ideas into cinematic visuals!")

with st.sidebar:
    st.header("🎨 Creative Controls")
    user_idea = st.text_area("Enter your core idea/story prompt:", "A lone astronaut discovers a glowing alien artifact on Mars, a sense of wonder and slight dread.", height=100)
    genre = st.selectbox("Select Genre:", ["Sci-Fi", "Fantasy", "Thriller", "Drama", "Comedy"], index=0)
    mood = st.selectbox("Select Mood:", ["Suspenseful", "Mysterious", "Epic", "Uplifting", "Dark"], index=0)
    generate_button = st.button("✨ Generate Cinematic Concept", type="primary")

# --- Session State for Storing Results ---
if 'story_scenes' not in st.session_state:
    st.session_state.story_scenes = None
if 'image_prompts' not in st.session_state:
    st.session_state.image_prompts = []
if 'generated_images_paths' not in st.session_state:
    st.session_state.generated_images_paths = []
if 'video_path' not in st.session_state:
    st.session_state.video_path = None

# --- Main Logic ---
if generate_button and user_idea:
    st.session_state.story_scenes = None # Reset previous results
    st.session_state.image_prompts = []
    st.session_state.generated_images_paths = []
    st.session_state.video_path = None

    with st.spinner("Phase 1: Gemini is drafting the script and scene breakdown... 📜"):
        story_prompt = create_story_breakdown_prompt(user_idea, genre, mood)
        st.session_state.story_scenes = gemini.generate_story_breakdown(story_prompt)

    if st.session_state.story_scenes:
        st.toast("Script and scene breakdown complete!", icon="✅")
        
        with st.spinner("Phase 2: Gemini is crafting visual prompts for keyframes... 🖼️"):
            for i, scene in enumerate(st.session_state.story_scenes):
                scene_desc_for_image = f"Scene {scene.get('scene_number', i+1)}: {scene.get('key_action', 'N/A')}. Setting: {scene.get('setting_description', 'N/A')}"
                img_gen_prompt_text = create_image_prompt_from_scene(scene_desc_for_image, scene.get('visual_style_suggestion', 'cinematic'))
                
                # Get the actual image description from Gemini
                # For this demo, we'll use the scene's key_action and setting as the "description" for the placeholder
                # In a real app, Gemini would generate a detailed image prompt here.
                # image_description_from_gemini = gemini.generate_image_prompt(img_gen_prompt_text)
                # For simplicity, we use parts of the scene data directly for the placeholder image text
                image_description_for_placeholder = f"Scene {scene.get('scene_number', i+1)}: {scene.get('key_action', 'N/A')}\nStyle: {scene.get('visual_style_suggestion', 'cinematic')}"
                
                if image_description_for_placeholder: # was image_description_from_gemini
                    st.session_state.image_prompts.append(image_description_for_placeholder)
                    # Simulate image generation by creating a placeholder
                    img_path = visualizer.create_placeholder_image(
                        image_description_for_placeholder, # use the generated desc
                        f"scene_{scene.get('scene_number', i+1)}_placeholder.png"
                    )
                    st.session_state.generated_images_paths.append(img_path)
                else:
                    st.warning(f"Could not generate image prompt/description for Scene {scene.get('scene_number', i+1)}.")
        st.toast("Visual prompts and placeholder images generated!", icon="🖼️")

        if st.session_state.generated_images_paths:
            with st.spinner("Phase 3: Assembling the cinematic video... 🎞️"):
                st.session_state.video_path = visualizer.create_video_from_images(
                    st.session_state.generated_images_paths,
                    output_filename="cinegen_output.mp4",
                    duration_per_image=3 # Show each image for 3 seconds
                )
            st.toast("Video assembled!", icon="🎬")
            st.balloons()
        else:
            st.error("No images were generated, cannot create video.")
    else:
        st.error("Failed to generate story breakdown from Gemini. Please check logs or try a different prompt.")

# --- Display Results ---
if st.session_state.story_scenes:
    st.header("📝 Generated Story Breakdown")
    for i, scene in enumerate(st.session_state.story_scenes):
        with st.expander(f"Scene {scene.get('scene_number', i+1)}: {scene.get('key_action', 'N/A')[:50]}...", expanded=i==0):
            st.markdown(f"**Setting:** {scene.get('setting_description', 'N/A')}")
            st.markdown(f"**Characters:** {', '.join(scene.get('characters_involved', []))}")
            st.markdown(f"**Key Action:** {scene.get('key_action', 'N/A')}")
            st.markdown(f"**Visual Style Suggestion:** {scene.get('visual_style_suggestion', 'N/A')}")
            st.markdown(f"**Camera Angle Suggestion:** {scene.get('camera_angle_suggestion', 'N/A')}")
            
            if i < len(st.session_state.image_prompts):
                 st.markdown(f"**Image Concept (Placeholder Text):**")
                 st.code(st.session_state.image_prompts[i])
            if i < len(st.session_state.generated_images_paths):
                st.image(st.session_state.generated_images_paths[i], caption=f"Visual Concept for Scene {scene.get('scene_number', i+1)}")


if st.session_state.video_path and os.path.exists(st.session_state.video_path):
    st.header("🎬 Generated Cinematic Concept Video")
    video_file = open(st.session_state.video_path, 'rb')
    video_bytes = video_file.read()
    st.video(video_bytes)
    st.markdown(f"Download your video: [{os.path.basename(st.session_state.video_path)}]({st.session_state.video_path})") # This download might not work directly on Spaces due to pathing, but video will play.
    # For download on Spaces, you might need to serve it differently or provide a direct link if Spaces keeps it.
    # A simpler way for Spaces is to just let the user right-click the video player and save.
else:
    if generate_button and not st.session_state.video_path and st.session_state.generated_images_paths:
        st.warning("Video generation might have failed or no images were available.")
    elif not generate_button :
        st.info("Enter your idea and click 'Generate Cinematic Concept' to begin!")