app.py

# 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_data,
    create_scene_regeneration_prompt,
    create_visual_regeneration_prompt
)
import os

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

# --- Global State Variables (Using session state for persistence) ---
# Load Gemini API Key
if 'GEMINI_API_KEY' not in st.session_state:
    try:
        st.session_state.GEMINI_API_KEY = st.secrets["GEMINI_API_KEY"]
    except KeyError: 
        if "GEMINI_API_KEY" in os.environ:
             st.session_state.GEMINI_API_KEY = os.environ["GEMINI_API_KEY"]
        else:
            st.error("GEMINI_API_KEY not found in secrets or environment variables. Please add it.")
            st.stop()
    except Exception as e: 
        st.error(f"Error loading Gemini API Key from secrets: {e}")
        st.stop()

# Initialize Gemini Handler
if 'gemini_handler' not in st.session_state:
    st.session_state.gemini_handler = GeminiHandler(api_key=st.session_state.GEMINI_API_KEY)

# Initialize Visual Engine and set OpenAI API Key
if 'visual_engine' not in st.session_state:
    st.session_state.visual_engine = VisualEngine(output_dir="temp_cinegen_media")
    openai_key_from_secrets = None
    try:
        if "OPENAI_API_KEY" in st.secrets:
            openai_key_from_secrets = st.secrets["OPENAI_API_KEY"]
    except AttributeError: 
        print("st.secrets not available (likely local dev without secrets.toml). Checking environment variables.")
    except Exception as e:
        print(f"Could not access st.secrets for OPENAI_API_KEY: {e}")

    if not openai_key_from_secrets and "OPENAI_API_KEY" in os.environ: 
        openai_key_from_secrets = os.environ["OPENAI_API_KEY"]
    
    if openai_key_from_secrets:
        st.session_state.visual_engine.set_openai_api_key(openai_key_from_secrets)
    else:
        st.session_state.visual_engine.set_openai_api_key(None) 
        st.sidebar.caption("OpenAI API Key for DALL-E not found. Visuals will be placeholders.")

# Story and generated content state
if 'story_scenes' not in st.session_state: 
    st.session_state.story_scenes = []
if 'scene_image_prompts' not in st.session_state: 
    st.session_state.scene_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
if 'character_definitions' not in st.session_state:
    st.session_state.character_definitions = {} 
if 'style_reference_description' not in st.session_state:
    st.session_state.style_reference_description = None

# --- Helper Functions ---
def initialize_new_story():
    st.session_state.story_scenes = []
    st.session_state.scene_image_prompts = []
    st.session_state.generated_images_paths = []
    st.session_state.video_path = None

def generate_visual_for_scene_wrapper(scene_index, scene_data, is_regeneration=False, version_count=1):
    scene_num_for_log = scene_data.get('scene_number', scene_index + 1)
    textual_image_prompt = ""
    if is_regeneration and scene_index < len(st.session_state.scene_image_prompts) and st.session_state.scene_image_prompts[scene_index]:
        textual_image_prompt = st.session_state.scene_image_prompts[scene_index]
    else:
        textual_image_prompt = create_image_prompt_from_scene_data(
            scene_data,
            st.session_state.character_definitions,
            st.session_state.style_reference_description
        )
    if not textual_image_prompt:
        return False
    if scene_index >= len(st.session_state.scene_image_prompts):
        while len(st.session_state.scene_image_prompts) <= scene_index:
             st.session_state.scene_image_prompts.append("")
    st.session_state.scene_image_prompts[scene_index] = textual_image_prompt
    image_filename = f"scene_{scene_num_for_log}_visual_v{version_count}.png"
    generated_image_path = st.session_state.visual_engine.generate_image_visual(
        textual_image_prompt, image_filename
    )
    while len(st.session_state.generated_images_paths) <= scene_index:
        st.session_state.generated_images_paths.append(None)
    if generated_image_path and os.path.exists(generated_image_path):
        st.session_state.generated_images_paths[scene_index] = generated_image_path
        return True
    else:
        st.session_state.generated_images_paths[scene_index] = None
        return False

# --- UI Sidebar ---
with st.sidebar:
    st.title("🎬 CineGen AI Pro")
    st.markdown("### Creative Controls")
    user_idea = st.text_area("Enter your core story idea:", "A lone astronaut discovers a glowing alien artifact on Mars, a sense of wonder and slight dread.", height=100, key="user_idea_input")
    genre = st.selectbox("Genre:", ["Sci-Fi", "Fantasy", "Noir", "Thriller", "Drama", "Horror", "Cyberpunk"], index=6, key="genre_select")
    mood = st.selectbox("Mood:", ["Suspenseful", "Mysterious", "Gritty", "Epic", "Dark", "Hopeful", "Wonder"], index=6, key="mood_select")
    num_scenes_val = st.slider("Number of Scenes:", 1, 5, 1, key="num_scenes_slider") # Default to 1 for faster testing

    if st.button("✨ Generate Full Story Concept", type="primary", key="generate_full_story_btn", use_container_width=True):
        initialize_new_story()
        if not user_idea.strip():
            st.warning("Please enter a story idea.")
        else:
            with st.status("Generating story...", expanded=True) as status_main:
                st.write("Phase 1: Gemini is drafting the script & scene breakdown... 📜")
                story_prompt_text = create_story_breakdown_prompt(user_idea, genre, mood, num_scenes_val)
                try:
                    st.session_state.story_scenes = st.session_state.gemini_handler.generate_story_breakdown(story_prompt_text)
                    status_main.update(label="Script breakdown complete! ✅", state="running", expanded=True)
                    num_actual_scenes = len(st.session_state.story_scenes)
                    st.session_state.scene_image_prompts = [""] * num_actual_scenes
                    st.session_state.generated_images_paths = [None] * num_actual_scenes
                except Exception as e:
                    status_main.update(label=f"Failed to generate story breakdown: {e}", state="error", expanded=True)
                    st.session_state.story_scenes = [] 
                    st.stop() 
                if st.session_state.story_scenes:
                    st.write("Phase 2: Generating initial visual concepts... 🖼️ (This may take time with DALL-E)")
                    success_count = 0
                    for i_loop, scene_data_loop_var in enumerate(st.session_state.story_scenes):
                        scene_num_log = scene_data_loop_var.get('scene_number', i_loop + 1)
                        st.write(f"Generating visual for Scene {scene_num_log}...")
                        if generate_visual_for_scene_wrapper(i_loop, scene_data_loop_var, version_count=1):
                            success_count +=1
                            st.write(f"Visual for Scene {scene_num_log} done.")
                        else:
                            st.write(f"Visual for Scene {scene_num_log} failed.")
                    if success_count == len(st.session_state.story_scenes):
                        status_main.update(label="All concepts generated successfully! 🎉", state="complete", expanded=False)
                    elif success_count > 0:
                         status_main.update(label=f"{success_count}/{len(st.session_state.story_scenes)} visuals generated.", state="warning", expanded=False)
                    else:
                        status_main.update(label="Visual concept generation failed for all scenes.", state="error", expanded=False)

    st.markdown("---")
    st.markdown("### Advanced Options")
    with st.expander("Character Consistency", expanded=False):
        char_name_input = st.text_input("Character Name (e.g., Eva)", key="char_name_adv_input")
        char_desc_input = st.text_area("Character Description (e.g., 'female, short red hair, green eyes, wearing a blue jumpsuit')", key="char_desc_adv_input", height=100)
        if st.button("Add/Update Character", key="add_char_adv_btn"):
            if char_name_input and char_desc_input:
                st.session_state.character_definitions[char_name_input.strip().lower()] = char_desc_input.strip()
                st.success(f"Character '{char_name_input.strip()}' defined.")
            else:
                st.warning("Please provide both name and description.")
        if st.session_state.character_definitions:
            st.caption("Defined Characters:")
            for char_key, desc_val in st.session_state.character_definitions.items():
                st.markdown(f"**{char_key.title()}:** _{desc_val}_")

    with st.expander("Style Transfer", expanded=False):
        style_ref_text = st.text_area("Describe Visual Style (e.g., 'impressionistic oil painting, vibrant colors')", key="style_text_adv_input", height=100)
        if st.button("Apply Textual Style", key="apply_style_adv_btn"):
            st.session_state.style_reference_description = style_ref_text.strip()
            st.success("Style reference applied. Re-generate visuals or full story to see changes.")

# --- Main Content Area ---
st.header("📝 Cinematic Storyboard")

if not st.session_state.story_scenes:
    st.info("Enter your idea in the sidebar and click 'Generate Full Story Concept' to begin.")
else:
    for i, scene_data_display in enumerate(st.session_state.story_scenes):
        scene_num_display = scene_data_display.get('scene_number', i + 1)
        action_summary = scene_data_display.get('key_action', f"scene{i}")
        key_part_raw = ''.join(e for e in action_summary if e.isalnum() or e.isspace())
        unique_key_part = key_part_raw.replace(" ", "_")[:20]

        st.subheader(f"Scene {scene_num_display}: {scene_data_display.get('emotional_beat', 'Untitled Scene')}")
        col1, col2 = st.columns([2, 3])

        with col1: 
            with st.expander("Scene Details", expanded=True):
                st.markdown(f"**Setting:** {scene_data_display.get('setting_description', 'N/A')}")
                st.markdown(f"**Characters:** {', '.join(scene_data_display.get('characters_involved', []))}")
                st.markdown(f"**Key Action:** {scene_data_display.get('key_action', 'N/A')}")
                st.markdown(f"**Dialogue Snippet:** `\"{scene_data_display.get('dialogue_snippet', '...')}\"`")
                st.markdown(f"**Visual Style:** {scene_data_display.get('visual_style_suggestion', 'N/A')}")
                st.markdown(f"**Camera:** {scene_data_display.get('camera_angle_suggestion', 'N/A')}")
                
                current_textual_prompt = st.session_state.scene_image_prompts[i] if i < len(st.session_state.scene_image_prompts) else None
                if current_textual_prompt:
                     with st.popover("View Image Prompt Text"): 
                        st.markdown(f"**Textual Prompt for Image Generation:**")
                        st.code(current_textual_prompt, language='text')
        with col2: 
            current_image_path = st.session_state.generated_images_paths[i] if i < len(st.session_state.generated_images_paths) else None
            if current_image_path and os.path.exists(current_image_path):
                # MODIFICATION HERE: Removed use_column_width=True
                st.image(current_image_path, caption=f"Visual Concept for Scene {scene_num_display}") 
            else:
                if st.session_state.story_scenes: 
                    st.caption("Visual for this scene is pending or failed to generate.")

            with st.popover(f"✏️ Edit Scene {scene_num_display} Script"): 
                feedback_script = st.text_area("Describe changes to script details:",
                    key=f"script_feedback_{unique_key_part}_{i}", height=100) 
                if st.button(f"🔄 Regenerate Scene {scene_num_display} Script", key=f"regen_script_btn_{unique_key_part}_{i}"): 
                    if feedback_script:
                        with st.status(f"Rewriting Scene {scene_num_display} script...", expanded=True) as status_script_regen:
                            regen_prompt = create_scene_regeneration_prompt(
                                scene_data_display, feedback_script, st.session_state.story_scenes
                            )
                            try:
                                updated_scene_data = st.session_state.gemini_handler.regenerate_scene_script_details(regen_prompt)
                                st.session_state.story_scenes[i] = updated_scene_data 
                                status_script_regen.update(label="Script updated! Regenerating visual...", state="running")
                                version = 1
                                if current_image_path:
                                    try: 
                                        base, _ = os.path.splitext(os.path.basename(current_image_path))
                                        if '_v' in base: version = int(base.split('_v')[-1]) + 1
                                    except : pass
                                if generate_visual_for_scene_wrapper(i, updated_scene_data, is_regeneration=True, version_count=version):
                                    status_script_regen.update(label="Scene script & visual updated! 🎉", state="complete", expanded=False)
                                else:
                                    status_script_regen.update(label="Script updated, but visual failed.", state="warning", expanded=False)
                                st.rerun()
                            except Exception as e:
                                status_script_regen.update(label=f"Error regenerating scene script: {e}", state="error")
                    else:
                        st.warning("Please provide feedback for script regeneration.")

            with st.popover(f"🎨 Edit Scene {scene_num_display} Visuals"): 
                prompt_for_edit_visuals = st.session_state.scene_image_prompts[i] if i < len(st.session_state.scene_image_prompts) else "No prompt text available."
                st.caption("Current Image Prompt Text:")
                st.code(prompt_for_edit_visuals, language='text')
                feedback_visual = st.text_area("Describe visual changes to apply to the prompt:",
                    key=f"visual_feedback_{unique_key_part}_{i}", height=100) 
                if st.button(f"🔄 Regenerate Scene {scene_num_display} Visuals", key=f"regen_visual_btn_{unique_key_part}_{i}"): 
                    if feedback_visual:
                        with st.status(f"Refining prompt & regenerating visual for Scene {scene_num_display}...", expanded=True) as status_visual_regen:
                            prompt_refinement_request = create_visual_regeneration_prompt(
                                prompt_for_edit_visuals, feedback_visual, scene_data_display 
                            )
                            try:
                                refined_textual_image_prompt = st.session_state.gemini_handler.regenerate_image_prompt_from_feedback(prompt_refinement_request)
                                st.session_state.scene_image_prompts[i] = refined_textual_image_prompt 
                                status_visual_regen.update(label="Image prompt refined! Regenerating visual...", state="running")
                                version = 1
                                if current_image_path:
                                    try:
                                        base, _ = os.path.splitext(os.path.basename(current_image_path))
                                        if '_v' in base: version = int(base.split('_v')[-1]) + 1
                                    except : pass
                                if generate_visual_for_scene_wrapper(i, scene_data_display, is_regeneration=True, version_count=version):
                                    status_visual_regen.update(label="Visual for Scene updated! 🎉", state="complete", expanded=False)
                                else:
                                     status_visual_regen.update(label="Prompt refined, but visual failed.", state="warning", expanded=False)
                                st.rerun()
                            except Exception as e:
                                status_visual_regen.update(label=f"Error refining prompt or regenerating visual: {e}", state="error")
                    else:
                        st.warning("Please provide feedback for visual regeneration.")
        st.markdown("---")
    
    if st.session_state.story_scenes and any(p for p in st.session_state.generated_images_paths if p is not None):
        if st.button("🎬 Assemble Enhanced Animatic", key="assemble_enhanced_video_btn", type="primary", use_container_width=True):
            with st.status("Assembling enhanced animatic video...", expanded=False) as status_video:
                image_data_for_video = []
                for idx, scene_info in enumerate(st.session_state.story_scenes):
                    img_path = st.session_state.generated_images_paths[idx] if idx < len(st.session_state.generated_images_paths) else None
                    if img_path and os.path.exists(img_path):
                        image_data_for_video.append({
                            'path': img_path,
                            'scene_num': scene_info.get('scene_number', idx + 1),
                            'key_action': scene_info.get('key_action', '') 
                        })
                if image_data_for_video:
                    st.session_state.video_path = st.session_state.visual_engine.create_video_from_images(
                        image_data_for_video, 
                        output_filename="cinegen_pro_animatic_enhanced.mp4",
                        duration_per_image=4, 
                        fps=24 
                    )
                    if st.session_state.video_path and os.path.exists(st.session_state.video_path):
                        status_video.update(label="Enhanced animatic assembled! 🎉", state="complete")
                        st.balloons()
                    else:
                        status_video.update(label="Enhanced video assembly failed. Check logs.", state="error")
                else:
                    status_video.update(label="No valid images to assemble video.", state="error")
    elif st.session_state.story_scenes:
         st.info("Generate visuals for scenes before assembling the video.")

    if st.session_state.video_path and os.path.exists(st.session_state.video_path):
        st.header("🎬 Generated Animatic Video")
        try:
            with open(st.session_state.video_path, 'rb') as video_file_obj:
                video_bytes_content = video_file_obj.read()
            st.video(video_bytes_content, format="video/mp4") 
            with open(st.session_state.video_path, "rb") as fp_download_video:
                st.download_button(
                    label="Download Enhanced Animatic", data=fp_download_video,
                    file_name=os.path.basename(st.session_state.video_path), mime="video/mp4",
                    use_container_width=True, key="download_video_btn" 
                )
        except Exception as e:
            st.error(f"Error displaying or preparing video for download: {e}")

# --- Footer ---
st.sidebar.markdown("---")
st.sidebar.caption("CineGen AI Pro | Powered by Gemini & Streamlit.")