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_cinematic_treatment_prompt,
    construct_dalle_prompt,
    create_narration_script_prompt_enhanced,
    create_scene_regeneration_prompt,
    create_visual_regeneration_prompt
)
import os
import logging

# --- Configuration & Initialization ---
st.set_page_config(page_title="CineGen AI Ultra+", layout="wide", initial_sidebar_state="expanded")
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# <<< ADDED/MODIFIED START >>>
# --- Global Definitions for New Features ---
SHOT_TYPES_OPTIONS = [
    "Director's Choice", "Establishing Shot", "Long Shot", "Full Shot",
    "Medium Long Shot (Cowboy)", "Medium Shot", "Medium Close-up",
    "Close-up", "Extreme Close-up", "Point of View (POV)",
    "Over the Shoulder", "Tracking Shot", "Dolly Zoom", "Crane Shot",
    "Aerial Shot", "Static Shot", "Dutch Angle", "Whip Pan"
]
DEFAULT_SCENE_DURATION_SECS = 5  # Default duration in seconds for each scene
DEFAULT_SHOT_TYPE = "Director's Choice"
# <<< ADDED/MODIFIED END >>>


# --- Global State Variables & API Key Setup ---
def load_api_key(key_name_streamlit, key_name_env, service_name):
    key = None; secrets_available = hasattr(st, 'secrets')
    try:
        if secrets_available and key_name_streamlit in st.secrets:
            key = st.secrets[key_name_streamlit]
            if key: logger.info(f"{service_name} API Key found in Streamlit secrets.")
    except Exception as e: logger.warning(f"Could not access st.secrets for {key_name_streamlit}: {e}")
    if not key and key_name_env in os.environ:
        key = os.environ[key_name_env]
        if key: logger.info(f"{service_name} API Key found in environment variable.")
    if not key: logger.warning(f"{service_name} API Key NOT FOUND. Related features may be disabled or use fallbacks.")
    return key

if 'services_initialized' not in st.session_state:
    logger.info("Initializing services and API keys for the first time this session...")
    st.session_state.GEMINI_API_KEY = load_api_key("GEMINI_API_KEY", "GEMINI_API_KEY", "Gemini")
    st.session_state.OPENAI_API_KEY = load_api_key("OPENAI_API_KEY", "OPENAI_API_KEY", "OpenAI/DALL-E")
    st.session_state.ELEVENLABS_API_KEY = load_api_key("ELEVENLABS_API_KEY", "ELEVENLABS_API_KEY", "ElevenLabs")
    st.session_state.PEXELS_API_KEY = load_api_key("PEXELS_API_KEY", "PEXELS_API_KEY", "Pexels")
    st.session_state.ELEVENLABS_VOICE_ID_CONFIG = load_api_key("ELEVENLABS_VOICE_ID", "ELEVENLABS_VOICE_ID", "ElevenLabs Voice ID")

    if not st.session_state.GEMINI_API_KEY:
        st.error("CRITICAL: Gemini API Key is essential and missing!"); logger.critical("Gemini API Key missing. Halting."); st.stop()

    try:
        st.session_state.gemini_handler = GeminiHandler(api_key=st.session_state.GEMINI_API_KEY)
        logger.info("GeminiHandler initialized successfully.")
    except Exception as e: st.error(f"Failed to init GeminiHandler: {e}"); logger.critical(f"GeminiHandler init failed: {e}", exc_info=True); st.stop()

    try:
        default_voice_id = "Rachel" # A common fallback if no secret is set
        configured_voice_id = st.session_state.ELEVENLABS_VOICE_ID_CONFIG or default_voice_id
        st.session_state.visual_engine = VisualEngine(
            output_dir="temp_cinegen_media",
            default_elevenlabs_voice_id=configured_voice_id # Pass it to __init__
        )
        st.session_state.visual_engine.set_openai_api_key(st.session_state.OPENAI_API_KEY)
        st.session_state.visual_engine.set_elevenlabs_api_key(st.session_state.ELEVENLABS_API_KEY, voice_id_from_secret=st.session_state.ELEVENLABS_VOICE_ID_CONFIG)
        st.session_state.visual_engine.set_pexels_api_key(st.session_state.PEXELS_API_KEY)
        logger.info("VisualEngine initialized and API keys set (or attempted).")
    except Exception as e:
        st.error(f"Failed to init VisualEngine or set its API keys: {e}"); logger.critical(f"VisualEngine init/key setting failed: {e}", exc_info=True)
        st.warning("VisualEngine critical setup issue. Some features will be disabled.")
    st.session_state.services_initialized = True; logger.info("Service initialization sequence complete.")

# Initialize other session state variables
for key, default_val in [
    ('story_treatment_scenes', []), ('scene_dalle_prompts', []), ('generated_visual_paths', []),
    ('video_path', None), ('character_definitions', {}), ('global_style_additions', ""),
    ('overall_narration_audio_path', None), ('narration_script_display', "")
]:
    if key not in st.session_state: st.session_state[key] = default_val

def initialize_new_project():
    st.session_state.story_treatment_scenes, st.session_state.scene_dalle_prompts, st.session_state.generated_visual_paths = [], [], []
    st.session_state.video_path, st.session_state.overall_narration_audio_path, st.session_state.narration_script_display = None, None, ""
    logger.info("New project initialized.")

def generate_visual_for_scene_core(scene_index, scene_data, version=1):
    dalle_prompt = construct_dalle_prompt(scene_data, st.session_state.character_definitions, st.session_state.global_style_additions)
    if not dalle_prompt: logger.error(f"DALL-E prompt construction failed for scene {scene_data.get('scene_number', scene_index+1)}"); return False
    while len(st.session_state.scene_dalle_prompts) <= scene_index: st.session_state.scene_dalle_prompts.append("")
    while len(st.session_state.generated_visual_paths) <= scene_index: st.session_state.generated_visual_paths.append(None)
    st.session_state.scene_dalle_prompts[scene_index] = dalle_prompt
    filename = f"scene_{scene_data.get('scene_number', scene_index+1)}_visual_v{version}.png"
    img_path = st.session_state.visual_engine.generate_image_visual(dalle_prompt, scene_data, filename)
    if img_path and os.path.exists(img_path):
        st.session_state.generated_visual_paths[scene_index] = img_path; logger.info(f"Visual generated for Scene {scene_data.get('scene_number', scene_index+1)}: {os.path.basename(img_path)}"); return True
    else:
        st.session_state.generated_visual_paths[scene_index] = None; logger.warning(f"Visual generation FAILED for Scene {scene_data.get('scene_number', scene_index+1)}. img_path was: {img_path}"); return False

# --- UI Sidebar ---
with st.sidebar:
    st.title("🎬 CineGen AI Ultra+")
    st.markdown("### Creative Seed")
    user_idea = st.text_area("Core Story Idea / Theme:", "A lone wanderer searches for a mythical oasis in a vast, post-apocalyptic desert, haunted by mirages and mechanical scavengers.", height=120, key="user_idea_main_v5")
    genre = st.selectbox("Primary Genre:", ["Cyberpunk", "Sci-Fi", "Fantasy", "Noir", "Thriller", "Western", "Post-Apocalyptic", "Historical Drama", "Surreal"], index=6, key="genre_main_v5")
    mood = st.selectbox("Overall Mood:", ["Hopeful yet Desperate", "Mysterious & Eerie", "Gritty & Tense", "Epic & Awe-Inspiring", "Melancholy & Reflective", "Whimsical & Lighthearted"], index=0, key="mood_main_v5")
    num_scenes = st.slider("Number of Key Scenes:", 1, 10, 3, key="num_scenes_main_v5") # MODIFIED: Increased max scenes
    creative_guidance_options = {"Standard Director": "standard", "Artistic Visionary": "more_artistic", "Experimental Storyteller": "experimental_narrative"}
    selected_creative_guidance_key = st.selectbox("AI Creative Director Style:", options=list(creative_guidance_options.keys()), key="creative_guidance_select_v5")
    actual_creative_guidance = creative_guidance_options[selected_creative_guidance_key]

    if st.button("🌌 Generate Cinematic Treatment", type="primary", key="generate_treatment_btn_v5", use_container_width=True):
        initialize_new_project()
        if not user_idea.strip(): st.warning("Please provide a story idea.")
        else:
            with st.status("AI Director is envisioning your masterpiece...", expanded=True) as status:
                try:
                    status.write("Phase 1: Gemini crafting cinematic treatment... 📜"); logger.info("Phase 1: Cinematic Treatment Gen.")
                    # Note: Consider updating create_cinematic_treatment_prompt to also ask Gemini
                    # for 'suggested_shot_type' and 'estimated_duration_secs' for each scene.
                    treatment_prompt = create_cinematic_treatment_prompt(user_idea, genre, mood, num_scenes, actual_creative_guidance)
                    treatment_result_json_raw = st.session_state.gemini_handler.generate_story_breakdown(treatment_prompt)
                    if not isinstance(treatment_result_json_raw, list) or not treatment_result_json_raw: raise ValueError("Gemini returned invalid scene list.")

                    # <<< ADDED/MODIFIED START >>>
                    # Process raw scenes and add default shot type and duration
                    processed_scenes = []
                    for scene_data_from_gemini in treatment_result_json_raw:
                        scene_data_from_gemini['shot_type'] = scene_data_from_gemini.get('suggested_shot_type', DEFAULT_SHOT_TYPE) # Use Gemini's suggestion if available
                        scene_data_from_gemini['scene_duration_secs'] = scene_data_from_gemini.get('estimated_duration_secs', DEFAULT_SCENE_DURATION_SECS) # Use Gemini's suggestion if available
                        processed_scenes.append(scene_data_from_gemini)
                    st.session_state.story_treatment_scenes = processed_scenes
                    # <<< ADDED/MODIFIED END >>>

                    num_gen_scenes = len(st.session_state.story_treatment_scenes)
                    st.session_state.scene_dalle_prompts = [""]*num_gen_scenes; st.session_state.generated_visual_paths = [None]*num_gen_scenes
                    logger.info(f"Phase 1 complete. {num_gen_scenes} scenes."); status.update(label="Treatment complete! ✅ Generating visuals...", state="running")

                    status.write("Phase 2: Creating visuals (DALL-E/Pexels)... 🖼️"); logger.info("Phase 2: Visual Gen.")
                    visual_successes = 0
                    for i, sc_data in enumerate(st.session_state.story_treatment_scenes):
                        sc_num_log = sc_data.get('scene_number', i+1)
                        status.write(f"  Visual for Scene {sc_num_log}..."); logger.info(f"  Processing visual for Scene {sc_num_log}.")
                        if generate_visual_for_scene_core(i, sc_data, version=1): visual_successes += 1

                    current_status_label_ph2 = "Visuals ready! "
                    next_step_state = "running"
                    if visual_successes == 0 and num_gen_scenes > 0:
                        logger.error("Visual gen failed all scenes."); current_status_label_ph2 = "Visual gen FAILED for all scenes."; next_step_state="error";
                        status.update(label=current_status_label_ph2, state=next_step_state, expanded=True); st.stop()
                    elif visual_successes < num_gen_scenes:
                        logger.warning(f"Visuals partial ({visual_successes}/{num_gen_scenes})."); current_status_label_ph2 = f"Visuals partially generated ({visual_successes}/{num_gen_scenes}). "
                    status.update(label=f"{current_status_label_ph2}Generating narration script...", state=next_step_state)
                    if next_step_state == "error": st.stop()

                    status.write("Phase 3: Generating narration script..."); logger.info("Phase 3: Narration Script Gen.")
                    voice_style_for_prompt = st.session_state.get("selected_voice_style_for_generation", "cinematic_trailer")
                    narr_prompt = create_narration_script_prompt_enhanced(st.session_state.story_treatment_scenes, mood, genre, voice_style_for_prompt)
                    st.session_state.narration_script_display = st.session_state.gemini_handler.generate_image_prompt(narr_prompt) # Assuming this is correct for narration text
                    logger.info("Narration script generated."); status.update(label="Narration script ready! Synthesizing voice...", state="running")

                    status.write("Phase 4: Synthesizing voice (ElevenLabs)... 🔊"); logger.info("Phase 4: Voice Synthesis.")
                    st.session_state.overall_narration_audio_path = st.session_state.visual_engine.generate_narration_audio(st.session_state.narration_script_display)

                    final_label = "All components ready! Storyboard below. 🚀"
                    final_state_val = "complete"
                    if not st.session_state.overall_narration_audio_path:
                        final_label = f"{current_status_label_ph2}Storyboard ready (Voiceover skipped or failed)."
                        logger.warning("Voiceover was skipped or failed.")
                    else: logger.info("Voiceover generated successfully.")
                    status.update(label=final_label, state=final_state_val, expanded=False)

                except ValueError as ve: logger.error(f"ValueError: {ve}", exc_info=True); status.update(label=f"Input or Gemini response error: {ve}", state="error", expanded=True);
                except Exception as e: logger.error(f"Unhandled Exception: {e}", exc_info=True); status.update(label=f"An unexpected error occurred: {e}", state="error", expanded=True);

    st.markdown("---"); st.markdown("### Fine-Tuning Options")
    with st.expander("Define Characters", expanded=False):
        char_name = st.text_input("Character Name", key="char_name_adv_ultra_v5"); char_desc = st.text_area("Visual Description", key="char_desc_adv_ultra_v5", height=100, placeholder="e.g., Jax: rugged male astronaut...")
        if st.button("Save Character", key="add_char_adv_ultra_v5"):
            if char_name and char_desc: st.session_state.character_definitions[char_name.strip().lower()] = char_desc.strip(); st.success(f"Char '{char_name.strip()}' saved.")
            else: st.warning("Name and description needed.")
        if st.session_state.character_definitions: st.caption("Current Characters:"); [st.markdown(f"**{k.title()}:** _{v}_") for k,v in st.session_state.character_definitions.items()]

    with st.expander("Global Style Overrides", expanded=False):
        presets = { "Default (Director's Choice)": "", "Hyper-Realistic Gritty Noir": "hyper-realistic gritty neo-noir, extreme detail, deep dynamic shadows, complex reflections on wet surfaces, cinematic film grain, desaturated palette with isolated vibrant neon accents (e.g. red, cyan), anamorphic lens distortion, atmospheric haze.", "Surreal Dreamscape Fantasy": "surreal dreamscape, epic fantasy elements, painterly with photorealistic details, impossible architecture, bioluminescent flora, otherworldly color palette (e.g., magenta skies, turquoise rivers), style of Roger Dean meets Zdzisław Beksiński.", "Vintage Analog Sci-Fi": "70s/80s analog sci-fi film aesthetic, tangible practical effects look, subtle light leaks, lens flares, warm filmic tones mixed with cool blues, detailed retro-futuristic technology with chunky buttons and CRT screens."}
        sel_preset = st.selectbox("Base Style Preset:", options=list(presets.keys()), key="style_preset_adv_ultra_v5")
        custom_kw = st.text_area("Additional Custom Style Keywords:", key="custom_style_adv_ultra_v5", height=80, placeholder="e.g., 'Dutch angle'")
        cur_style = st.session_state.global_style_additions
        if st.button("Apply Global Styles", key="apply_styles_adv_ultra_v5"):
            final_s = presets[sel_preset];
            if custom_kw.strip(): final_s = f"{final_s}, {custom_kw.strip()}" if final_s else custom_kw.strip()
            st.session_state.global_style_additions = final_s.strip(); cur_style = final_s.strip()
            if cur_style: st.success("Global styles applied!")
            else: st.info("Global style additions cleared.")
        if cur_style: st.caption(f"Active global styles: \"{cur_style}\"")

    with st.expander("Voice & Narration Style", expanded=False):
        default_voice_from_engine = "Rachel"
        if hasattr(st.session_state, 'visual_engine') and st.session_state.visual_engine:
            default_voice_from_engine = st.session_state.visual_engine.elevenlabs_voice_id

        user_voice_id_override = st.text_input(
            "ElevenLabs Voice ID (optional override):",
            value=default_voice_from_engine,
            key="el_voice_id_override_v5",
            help=f"Defaulting to '{default_voice_from_engine}' from secrets/config. Enter a specific Voice ID from your ElevenLabs account to override."
        )
        prompt_v_styles = {"Cinematic Trailer": "cinematic_trailer", "Neutral Documentary": "documentary_neutral", "Character Introspection": "introspective_character"}
        sel_prompt_v_style_key = st.selectbox("Narration Script Style:", list(prompt_v_styles.keys()), key="narr_style_sel_v5", index=0)

        if st.button("Set Narrator Voice & Style", key="set_voice_btn_ultra_v5"):
            final_voice_id_to_use = user_voice_id_override.strip()
            if not final_voice_id_to_use:
                final_voice_id_to_use = st.session_state.get("ELEVENLABS_VOICE_ID_CONFIG", "Rachel")

            if hasattr(st.session_state, 'visual_engine'):
                st.session_state.visual_engine.elevenlabs_voice_id = final_voice_id_to_use
            st.session_state.selected_voice_style_for_generation = prompt_v_styles[sel_prompt_v_style_key]
            st.success(f"Narrator Voice ID set to: {final_voice_id_to_use}. Script Style: {sel_prompt_v_style_key}")
            logger.info(f"User updated ElevenLabs Voice ID to: {final_voice_id_to_use}, Script Style: {sel_prompt_v_style_key}")

# --- Main Content Area ---
st.header("🎬 Cinematic Storyboard & Treatment")
if st.session_state.narration_script_display:
    with st.expander("📜 View Full Narration Script", expanded=False): st.markdown(f"> _{st.session_state.narration_script_display}_")

if not st.session_state.story_treatment_scenes: st.info("Use the sidebar to generate your cinematic treatment.")
else:
    for i_main, scene_content_display in enumerate(st.session_state.story_treatment_scenes):
        scene_n = scene_content_display.get('scene_number', i_main + 1); scene_t = scene_content_display.get('scene_title', 'Untitled')
        key_base = f"s{scene_n}_{''.join(filter(str.isalnum, scene_t[:10]))}_v5_{i_main}" # Ensure unique key base with index
        if "director_note" in scene_content_display and scene_content_display['director_note']: st.info(f"🎬 Director Note S{scene_n}: {scene_content_display['director_note']}")
        st.subheader(f"SCENE {scene_n}: {scene_t.upper()}"); col_d, col_v = st.columns([0.45, 0.55])
        with col_d:
            with st.expander("📝 Scene Treatment & Controls", expanded=True): # MODIFIED Expander Title
                st.markdown(f"**Beat:** {scene_content_display.get('emotional_beat', 'N/A')}"); st.markdown(f"**Setting:** {scene_content_display.get('setting_description', 'N/A')}"); st.markdown(f"**Chars:** {', '.join(scene_content_display.get('characters_involved', ['N/A']))}"); st.markdown(f"**Focus Moment:** _{scene_content_display.get('character_focus_moment', 'N/A')}_"); st.markdown(f"**Plot Beat:** {scene_content_display.get('key_plot_beat', 'N/A')}"); st.markdown(f"**Dialogue Hook:** `\"{scene_content_display.get('suggested_dialogue_hook', '...')}\"`"); st.markdown("---"); st.markdown(f"**Dir. Visual Style:** _{scene_content_display.get('PROACTIVE_visual_style_감독', 'N/A')}_"); st.markdown(f"**Dir. Camera:** _{scene_content_display.get('PROACTIVE_camera_work_감독', 'N/A')}_"); st.markdown(f"**Dir. Sound:** _{scene_content_display.get('PROACTIVE_sound_design_감독', 'N/A')}_")

                # <<< ADDED/MODIFIED START >>>
                st.markdown("---")
                st.markdown("##### Shot & Pacing Controls")

                # Shot Type Selection
                current_shot_type = st.session_state.story_treatment_scenes[i_main].get('shot_type', DEFAULT_SHOT_TYPE)
                try:
                    shot_type_index = SHOT_TYPES_OPTIONS.index(current_shot_type)
                except ValueError:
                    shot_type_index = SHOT_TYPES_OPTIONS.index(DEFAULT_SHOT_TYPE) # Fallback if value is somehow invalid

                new_shot_type = st.selectbox(
                    "Dominant Shot Type:",
                    options=SHOT_TYPES_OPTIONS,
                    index=shot_type_index,
                    key=f"shot_type_widget_{key_base}",
                    help="Suggests the primary camera shot for this scene. Influences visual generation style if DALL-E prompt is ever made aware of it."
                )
                if new_shot_type != st.session_state.story_treatment_scenes[i_main]['shot_type']:
                    st.session_state.story_treatment_scenes[i_main]['shot_type'] = new_shot_type
                    # No st.rerun() needed unless other UI elements immediately depend on this change

                # Scene Duration Control
                current_duration = st.session_state.story_treatment_scenes[i_main].get('scene_duration_secs', DEFAULT_SCENE_DURATION_SECS)
                new_duration = st.number_input(
                    "Scene Duration (seconds):",
                    min_value=1,
                    max_value=300, # Max 5 minutes per scene image
                    value=current_duration,
                    step=1,
                    key=f"duration_widget_{key_base}",
                    help="Approximate duration this scene's visual will be shown in the final animatic."
                )
                if new_duration != st.session_state.story_treatment_scenes[i_main]['scene_duration_secs']:
                    st.session_state.story_treatment_scenes[i_main]['scene_duration_secs'] = new_duration
                    # No st.rerun() needed unless other UI elements immediately depend on this change
                st.markdown("---")
                # <<< ADDED/MODIFIED END >>>

                cur_d_prompt = st.session_state.scene_dalle_prompts[i_main] if i_main < len(st.session_state.scene_dalle_prompts) else None
                if cur_d_prompt:
                     with st.popover("👁️ DALL-E Prompt"): st.markdown(f"**DALL-E Prompt:**"); st.code(cur_d_prompt, language='text')
                pexels_q = scene_content_display.get('pexels_search_query_감독', None)
                if pexels_q: st.caption(f"Pexels Fallback Query: `{pexels_q}`")

        with col_v:
            cur_img_p = st.session_state.generated_visual_paths[i_main] if i_main < len(st.session_state.generated_visual_paths) else None
            if cur_img_p and os.path.exists(cur_img_p): st.image(cur_img_p, caption=f"Scene {scene_n}: {scene_t}")
            else:
                if st.session_state.story_treatment_scenes: st.caption("Visual pending/failed.")

            with st.popover(f"✏️ Edit Scene {scene_n} Treatment"):
                fb_script = st.text_area("Changes to treatment:", key=f"treat_fb_{key_base}", height=150)
                if st.button(f"🔄 Update Scene {scene_n} Treatment", key=f"regen_treat_btn_{key_base}"):
                    if fb_script:
                        with st.status(f"Updating Scene {scene_n}...", expanded=True) as s_treat_regen:
                            prompt_text = create_scene_regeneration_prompt(scene_content_display, fb_script, st.session_state.story_treatment_scenes)
                            try:
                                updated_sc_data_raw = st.session_state.gemini_handler.regenerate_scene_script_details(prompt_text) # This is just the script part

                                # <<< ADDED/MODIFIED START >>>
                                # Preserve user-set shot type and duration
                                original_shot_type = st.session_state.story_treatment_scenes[i_main].get('shot_type', DEFAULT_SHOT_TYPE)
                                original_duration = st.session_state.story_treatment_scenes[i_main].get('scene_duration_secs', DEFAULT_SCENE_DURATION_SECS)

                                # Merge Gemini's script updates with existing scene data, preserving our custom fields
                                updated_sc_data = {**st.session_state.story_treatment_scenes[i_main], **updated_sc_data_raw}
                                updated_sc_data['shot_type'] = updated_sc_data_raw.get('suggested_shot_type', original_shot_type) # If Gemini suggests it, use it, else keep old
                                updated_sc_data['scene_duration_secs'] = updated_sc_data_raw.get('estimated_duration_secs', original_duration) # Same for duration

                                st.session_state.story_treatment_scenes[i_main] = updated_sc_data
                                # <<< ADDED/MODIFIED END >>>

                                s_treat_regen.update(label="Treatment updated! Regenerating visual...", state="running")
                                v_num = 1
                                if cur_img_p and os.path.exists(cur_img_p):
                                    try: b,_=os.path.splitext(os.path.basename(cur_img_p)); v_num = int(b.split('_v')[-1])+1 if '_v' in b else 2
                                    except (ValueError, IndexError, TypeError): v_num = 2
                                else: v_num = 1
                                if generate_visual_for_scene_core(i_main, updated_sc_data, version=v_num): s_treat_regen.update(label="Treatment & Visual Updated! 🎉", state="complete", expanded=False)
                                else: s_treat_regen.update(label="Treatment updated, visual failed.", state="complete", expanded=False)
                                st.rerun()
                            except Exception as e_regen: s_treat_regen.update(label=f"Error: {e_regen}", state="error"); logger.error(f"Scene treatment regen error: {e_regen}", exc_info=True)
                    else: st.warning("Please provide feedback.")

            with st.popover(f"🎨 Edit Scene {scene_n} Visual Prompt"):
                d_prompt_edit = st.session_state.scene_dalle_prompts[i_main] if i_main < len(st.session_state.scene_dalle_prompts) else "No DALL-E prompt."
                st.caption("Current DALL-E Prompt:"); st.code(d_prompt_edit, language='text')
                fb_visual = st.text_area("Changes for DALL-E prompt:", key=f"visual_fb_{key_base}", height=150)
                if st.button(f"🔄 Update Scene {scene_n} Visual", key=f"regen_visual_btn_{key_base}"):
                    if fb_visual:
                        with st.status(f"Refining prompt & visual for Scene {scene_n}...", expanded=True) as s_visual_regen:
                            ref_req_prompt = create_visual_regeneration_prompt(d_prompt_edit, fb_visual, scene_content_display,
                                st.session_state.character_definitions, st.session_state.global_style_additions)
                            try:
                                refined_d_prompt = st.session_state.gemini_handler.generate_image_prompt(ref_req_prompt)
                                st.session_state.scene_dalle_prompts[i_main] = refined_d_prompt
                                s_visual_regen.update(label="DALL-E prompt refined! Regenerating visual...", state="running")
                                v_num = 1
                                if cur_img_p and os.path.exists(cur_img_p):
                                    try: b,_=os.path.splitext(os.path.basename(cur_img_p)); v_num = int(b.split('_v')[-1])+1 if '_v' in b else 2
                                    except (ValueError, IndexError, TypeError): v_num=2
                                else: v_num = 1
                                # Pass the current scene_content_display, as DALL-E prompt is separate from treatment text
                                if generate_visual_for_scene_core(i_main, st.session_state.story_treatment_scenes[i_main], version=v_num): s_visual_regen.update(label="Visual Updated! 🎉", state="complete", expanded=False)
                                else: s_visual_regen.update(label="Prompt refined, visual failed.", state="complete", expanded=False)
                                st.rerun()
                            except Exception as e_regen_vis: s_visual_regen.update(label=f"Error: {e_regen_vis}", state="error"); logger.error(f"Visual prompt regen error: {e_regen_vis}", exc_info=True)
                    else: st.warning("Please provide feedback.")
        st.markdown("---")

    if st.session_state.story_treatment_scenes and any(p for p in st.session_state.generated_visual_paths if p is not None):
        if st.button("🎬 Assemble Narrated Cinematic Animatic", key="assemble_ultra_video_btn_v5", type="primary", use_container_width=True):
            with st.status("Assembling Ultra Animatic...", expanded=True) as status_vid:
                img_data_vid = []
                for i_v, sc_c in enumerate(st.session_state.story_treatment_scenes):
                    img_p_v = st.session_state.generated_visual_paths[i_v] if i_v < len(st.session_state.generated_visual_paths) else None
                    if img_p_v and os.path.exists(img_p_v):
                        # <<< ADDED/MODIFIED START >>>
                        scene_duration = sc_c.get('scene_duration_secs', DEFAULT_SCENE_DURATION_SECS)
                        img_data_vid.append({
                            'path': img_p_v,
                            'scene_num': sc_c.get('scene_number', i_v + 1),
                            'key_action': sc_c.get('key_plot_beat', ''),
                            'duration': scene_duration # Use per-scene duration
                        })
                        status_vid.write(f"Adding Scene {sc_c.get('scene_number', i_v + 1)} (Duration: {scene_duration}s).")
                        # <<< ADDED/MODIFIED END >>>
                    else:
                        logger.warning(f"Skipping Scene {sc_c.get('scene_number', i_v + 1)} for video assembly: No valid visual path.")


                if img_data_vid:
                    status_vid.write("Calling video engine...");
                    # <<< ADDED/MODIFIED START >>>
                    # IMPORTANT: Assumes visual_engine.create_video_from_images is updated to accept
                    # images_data (list of dicts with 'path' and 'duration') and no longer needs duration_per_image.
                    st.session_state.video_path = st.session_state.visual_engine.create_video_from_images(
                        images_data=img_data_vid, # Pass the list of dicts
                        overall_narration_path=st.session_state.overall_narration_audio_path,
                        output_filename="cinegen_ultra_animatic.mp4",
                        fps=24 # Keep FPS or make it configurable
                    )
                    # <<< ADDED/MODIFIED END >>>
                    if st.session_state.video_path and os.path.exists(st.session_state.video_path): status_vid.update(label="Ultra animatic assembled! 🎉", state="complete", expanded=False); st.balloons()
                    else: status_vid.update(label="Video assembly failed. Check logs.", state="error", expanded=False); logger.error("Video assembly returned None or file does not exist.")
                else: status_vid.update(label="No valid images with duration for video.", state="error", expanded=False); logger.warning("No valid images found for video assembly.")
    elif st.session_state.story_treatment_scenes: st.info("Generate visuals before assembling video.")

    if st.session_state.video_path and os.path.exists(st.session_state.video_path):
        st.header("🎬 Generated Cinematic Animatic");
        try:
            with open(st.session_state.video_path, 'rb') as vf_obj: video_bytes = vf_obj.read()
            st.video(video_bytes, format="video/mp4")
            with open(st.session_state.video_path, "rb") as fp_dl:
                st.download_button(label="Download Ultra Animatic", data=fp_dl, file_name=os.path.basename(st.session_state.video_path), mime="video/mp4", use_container_width=True, key="download_ultra_video_btn_v5" )
        except Exception as e: st.error(f"Error displaying video: {e}"); logger.error(f"Error displaying video: {e}", exc_info=True)

# --- Footer ---
st.sidebar.markdown("---"); st.sidebar.caption("CineGen AI Ultra+ | Visionary Cinematic Pre-Production")