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,
    construct_text_to_video_prompt, # Import new function
    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__)

# --- 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_SHOT_TYPE = "Director's Choice"
ASSET_TYPE_OPTIONS = ["Auto (Director's Choice)", "Image", "Video Clip"] # For user selection


# --- 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")
    st.session_state.RUNWAY_API_KEY = load_api_key("RUNWAY_API_KEY", "RUNWAY_API_KEY", "RunwayML") # Load Runway Key

    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"
        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
        )
        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)
        st.session_state.visual_engine.set_runway_api_key(st.session_state.RUNWAY_API_KEY) # Set Runway 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
# <<< MODIFIED START >>> : Renamed generated_visual_paths to generated_scene_assets
for key, default_val in [
    ('story_treatment_scenes', []), ('scene_prompts', []), ('generated_scene_assets', []), # Stores dicts: {'path':..., 'type':...}
    ('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_prompts = [] # Stores DALL-E or Text-to-Video prompts
    st.session_state.generated_scene_assets = [] # Stores dicts {'path': ..., 'type': ..., 'error': ...}
    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.")
# <<< MODIFIED END >>>

# <<< MODIFIED START >>> : Updated function to use generate_scene_asset
def generate_asset_for_scene_core(scene_index, scene_data, version=1, user_selected_asset_type="Auto (Director's Choice)"):
    """
    Generates a visual asset (image or video clip) for a scene.
    Returns True on success, False on failure.
    """
    # Determine asset type: user override > Gemini suggestion > default to image
    final_asset_type_decision = "image" # Default
    gemini_suggested_type = scene_data.get('suggested_asset_type_감독', 'image').lower()

    if user_selected_asset_type == "Image":
        final_asset_type_decision = "image"
    elif user_selected_asset_type == "Video Clip":
        final_asset_type_decision = "video_clip"
    elif user_selected_asset_type == "Auto (Director's Choice)":
        final_asset_type_decision = gemini_suggested_type if gemini_suggested_type == "video_clip" else "image"

    generate_as_video = (final_asset_type_decision == "video_clip")
    prompt_text_for_visual = ""

    if generate_as_video:
        # Construct prompt for text-to-video (e.g., RunwayML)
        prompt_text_for_visual = construct_text_to_video_prompt(scene_data, st.session_state.character_definitions, st.session_state.global_style_additions)
        # Note: seed_image_path could be an enhancement if DALL-E image is generated first
    else:
        # Construct prompt for DALL-E (image)
        prompt_text_for_visual = construct_dalle_prompt(scene_data, st.session_state.character_definitions, st.session_state.global_style_additions)

    if not prompt_text_for_visual:
        logger.error(f"Visual prompt construction failed for scene {scene_data.get('scene_number', scene_index+1)} (Type: {final_asset_type_decision})")
        return False

    # Ensure session state lists are long enough
    while len(st.session_state.scene_prompts) <= scene_index: st.session_state.scene_prompts.append("")
    while len(st.session_state.generated_scene_assets) <= scene_index: st.session_state.generated_scene_assets.append(None)

    st.session_state.scene_prompts[scene_index] = prompt_text_for_visual

    # Filename base (extension will be added by visual_engine)
    filename_base = f"scene_{scene_data.get('scene_number', scene_index+1)}_asset_v{version}"
    runway_duration = scene_data.get('video_clip_duration_estimate_secs_감독', DEFAULT_SCENE_DURATION_SECS)
    if runway_duration <= 0 : runway_duration = DEFAULT_SCENE_DURATION_SECS # Ensure positive duration

    asset_result = st.session_state.visual_engine.generate_scene_asset(
        image_prompt_text=prompt_text_for_visual, # This is generic, used for DALL-E or T2V
        scene_data=scene_data,
        scene_identifier_filename_base=filename_base,
        generate_as_video_clip=generate_as_video,
        runway_target_duration=runway_duration
        # input_image_for_runway=None # TODO: Could be an enhancement
    )

    st.session_state.generated_scene_assets[scene_index] = asset_result # Store the whole dict

    if asset_result and not asset_result['error'] and asset_result.get('path') and os.path.exists(asset_result['path']):
        logger.info(f"Asset ({asset_result.get('type')}) generated for Scene {scene_data.get('scene_number', scene_index+1)}: {os.path.basename(asset_result['path'])}")
        return True
    else:
        err_msg = asset_result.get('error_message', 'Unknown error') if asset_result else 'Asset result is None'
        logger.warning(f"Asset generation FAILED for Scene {scene_data.get('scene_number', scene_index+1)}. Type attempted: {final_asset_type_decision}. Path was: {asset_result.get('path') if asset_result else 'N/A'}. Error: {err_msg}")
        # Store a failure state
        st.session_state.generated_scene_assets[scene_index] = {'path': None, 'type': 'none', 'error': True, 'error_message': err_msg, 'prompt_used': prompt_text_for_visual}
        return False
# <<< MODIFIED END >>>

# --- 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, 2, key="num_scenes_main_v5")
    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.")
                    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) # Expect list of dicts
                    if not isinstance(treatment_result_json_raw, list) or not treatment_result_json_raw: raise ValueError("Gemini returned invalid scene list format.")

                    processed_scenes = []
                    for scene_data_from_gemini in treatment_result_json_raw:
                        scene_data_from_gemini['user_shot_type'] = scene_data_from_gemini.get('PROACTIVE_camera_work_감독', DEFAULT_SHOT_TYPE) # Default from Gemini's suggestion
                        scene_data_from_gemini['user_scene_duration_secs'] = scene_data_from_gemini.get('video_clip_duration_estimate_secs_감독', DEFAULT_SCENE_DURATION_SECS)
                        if scene_data_from_gemini['user_scene_duration_secs'] <=0: scene_data_from_gemini['user_scene_duration_secs'] = DEFAULT_SCENE_DURATION_SECS
                        scene_data_from_gemini['user_selected_asset_type'] = "Auto (Director's Choice)" # Default for UI
                        processed_scenes.append(scene_data_from_gemini)
                    st.session_state.story_treatment_scenes = processed_scenes

                    num_gen_scenes = len(st.session_state.story_treatment_scenes)
                    # <<< MODIFIED START >>>
                    st.session_state.scene_prompts = [""]*num_gen_scenes
                    st.session_state.generated_scene_assets = [None]*num_gen_scenes # Initialize list for asset dicts
                    # <<< MODIFIED END >>>
                    logger.info(f"Phase 1 complete. {num_gen_scenes} scenes."); status.update(label="Treatment complete! ✅ Generating visuals...", state="running")

                    status.write("Phase 2: Creating visual assets (Image/Video)... 🖼️🎬"); logger.info("Phase 2: Visual Asset 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"  Asset for Scene {sc_num_log}..."); logger.info(f"  Processing asset for Scene {sc_num_log}.")
                        # <<< MODIFIED START >>> : Calling new function
                        if generate_asset_for_scene_core(i, sc_data, version=1): # Default to 'Auto' asset type for initial gen
                            visual_successes += 1
                        # <<< MODIFIED END >>>

                    current_status_label_ph2 = "Visual assets ready! "
                    next_step_state = "running"
                    if visual_successes == 0 and num_gen_scenes > 0:
                        logger.error("Visual asset gen failed for all scenes."); current_status_label_ph2 = "Asset 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"Assets partially generated ({visual_successes}/{num_gen_scenes})."); current_status_label_ph2 = f"Assets 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) # This generates a string
                    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")
    # ... (Character, Global Style, Voice expanders - no changes needed here for this fix) ...
    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}"
        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: # Treatment and Controls Column
            with st.expander("📝 Scene Treatment & Controls", expanded=True):
                # Display scene textual details (emotional_beat, setting, etc.)
                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')}_")
                st.markdown("---")
                st.markdown("##### Shot, Pacing & Asset Controls")

                # User Shot Type (Camera Angle)
                current_shot_type = st.session_state.story_treatment_scenes[i_main].get('user_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)
                new_shot_type = st.selectbox("Dominant Shot Type:", options=SHOT_TYPES_OPTIONS, index=shot_type_index, key=f"shot_type_widget_{key_base}")
                if new_shot_type != current_shot_type:
                    st.session_state.story_treatment_scenes[i_main]['user_shot_type'] = new_shot_type
                    # Consider if a re-run is needed or if DALL-E prompt should be updated based on this

                # User Scene Duration
                current_duration = st.session_state.story_treatment_scenes[i_main].get('user_scene_duration_secs', DEFAULT_SCENE_DURATION_SECS)
                new_duration = st.number_input("Scene Duration (seconds):", min_value=1, max_value=300, value=current_duration, step=1, key=f"duration_widget_{key_base}")
                if new_duration != current_duration:
                    st.session_state.story_treatment_scenes[i_main]['user_scene_duration_secs'] = new_duration

                # <<< MODIFIED START >>> : User Asset Type Selection
                current_user_asset_type = st.session_state.story_treatment_scenes[i_main].get('user_selected_asset_type', "Auto (Director's Choice)")
                try: asset_type_idx = ASSET_TYPE_OPTIONS.index(current_user_asset_type)
                except ValueError: asset_type_idx = 0 # Default to Auto
                new_user_asset_type = st.selectbox("Asset Type Override:", ASSET_TYPE_OPTIONS, index=asset_type_idx, key=f"asset_type_sel_{key_base}",
                                                   help="Choose 'Image' or 'Video Clip'. 'Auto' uses Gemini's suggestion.")
                if new_user_asset_type != current_user_asset_type:
                    st.session_state.story_treatment_scenes[i_main]['user_selected_asset_type'] = new_user_asset_type
                    # This change will be picked up by regeneration buttons
                # <<< MODIFIED END >>>

                st.markdown("---")
                # Display generated prompt for the asset
                current_prompt_for_asset = st.session_state.scene_prompts[i_main] if i_main < len(st.session_state.scene_prompts) else None
                if current_prompt_for_asset:
                     with st.popover("👁️ View Asset Generation Prompt"):
                        st.markdown(f"**Prompt used for current asset:**"); st.code(current_prompt_for_asset, 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: # Visuals Column
            # <<< MODIFIED START >>> : Display logic for different asset types
            current_asset_data = st.session_state.generated_scene_assets[i_main] if i_main < len(st.session_state.generated_scene_assets) else None
            if current_asset_data and not current_asset_data.get('error') and current_asset_data.get('path') and os.path.exists(current_asset_data['path']):
                asset_path = current_asset_data['path']
                asset_type = current_asset_data.get('type', 'image') # Default to image if type missing
                if asset_type == 'image':
                    st.image(asset_path, caption=f"Scene {scene_n} ({asset_type}): {scene_t}")
                elif asset_type == 'video':
                    try:
                        with open(asset_path, 'rb') as vf: video_bytes = vf.read()
                        st.video(video_bytes, format="video/mp4", start_time=0)
                        st.caption(f"Scene {scene_n} ({asset_type}): {scene_t}")
                    except Exception as e_vid:
                        st.error(f"Error displaying video {asset_path}: {e_vid}")
                        logger.error(f"Error displaying video {asset_path}: {e_vid}", exc_info=True)
                else:
                    st.warning(f"Unknown asset type '{asset_type}' for Scene {scene_n}.")
            else: # No asset, or error during generation
                if st.session_state.story_treatment_scenes: # Check if treatment exists
                    error_msg = current_asset_data.get('error_message', 'Visual pending or failed.') if current_asset_data else 'Visual pending or failed.'
                    st.caption(error_msg)
            # <<< MODIFIED END >>>

            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} Treatment & Asset...", expanded=True) as s_treat_regen:
                            # Preserve user's shot type, duration, and asset type choices
                            user_shot_type = st.session_state.story_treatment_scenes[i_main]['user_shot_type']
                            user_duration = st.session_state.story_treatment_scenes[i_main]['user_scene_duration_secs']
                            user_asset_type_choice = st.session_state.story_treatment_scenes[i_main]['user_selected_asset_type']

                            prompt_text = create_scene_regeneration_prompt(scene_content_display, fb_script, st.session_state.story_treatment_scenes)
                            try:
                                updated_sc_data_from_gemini = st.session_state.gemini_handler.regenerate_scene_script_details(prompt_text)
                                # Merge, but prioritize user's UI choices for duration/shot/asset type
                                updated_sc_data = {**updated_sc_data_from_gemini} # Start with Gemini's new script
                                updated_sc_data['user_shot_type'] = user_shot_type
                                updated_sc_data['user_scene_duration_secs'] = user_duration
                                updated_sc_data['user_selected_asset_type'] = user_asset_type_choice
                                # Gemini might re-suggest asset type/duration, but user's direct settings take precedence for next gen
                                # We can log if Gemini's suggestion differs from user's explicit choice.
                                if updated_sc_data.get('suggested_asset_type_감독') != user_asset_type_choice and user_asset_type_choice != "Auto (Director's Choice)":
                                    logger.info(f"Scene {scene_n}: User asset choice '{user_asset_type_choice}' overrides Gemini suggestion '{updated_sc_data.get('suggested_asset_type_감독')}'.")


                                st.session_state.story_treatment_scenes[i_main] = updated_sc_data
                                s_treat_regen.update(label="Treatment updated! Regenerating asset...", state="running")

                                v_num = 1
                                if current_asset_data and current_asset_data.get('path') and os.path.exists(current_asset_data['path']):
                                    try: b,_=os.path.splitext(os.path.basename(current_asset_data['path'])); v_num = int(b.split('_v')[-1])+1 if '_v' in b else 2
                                    except: v_num = 2
                                else: v_num = 1
                                # <<< MODIFIED START >>> : Call new function, pass user_selected_asset_type
                                if generate_asset_for_scene_core(i_main, updated_sc_data, version=v_num, user_selected_asset_type=user_asset_type_choice):
                                    s_treat_regen.update(label="Treatment & Asset Updated! 🎉", state="complete", expanded=False)
                                else: s_treat_regen.update(label="Treatment updated, asset failed.", state="complete", expanded=False)
                                # <<< MODIFIED END >>>
                                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"):
                prompt_to_edit = st.session_state.scene_prompts[i_main] if i_main < len(st.session_state.scene_prompts) else "No prompt generated yet."
                st.caption("Current Asset Generation Prompt:"); st.code(prompt_to_edit, language='text')
                fb_visual = st.text_area("Changes for asset generation prompt:", key=f"visual_fb_{key_base}", height=150)
                if st.button(f"🔄 Update Scene {scene_n} Asset", key=f"regen_visual_btn_{key_base}"):
                    if fb_visual:
                        with st.status(f"Refining prompt & asset for Scene {scene_n}...", expanded=True) as s_visual_regen:
                            user_asset_type_choice = st.session_state.story_treatment_scenes[i_main]['user_selected_asset_type']
                            is_video_prompt = (user_asset_type_choice == "Video Clip") or \
                                              (user_asset_type_choice == "Auto (Director's Choice)" and \
                                               scene_content_display.get('suggested_asset_type_감독') == 'video_clip')

                            # Note: Visual regeneration prompt is primarily for DALL-E (images).
                            # For video, we might need a different refinement strategy or just regenerate with the same prompt construction.
                            # For simplicity here, if it's a video, we'll regenerate the prompt using standard construction.
                            # If it's an image, we use Gemini to refine the DALL-E prompt.
                            new_asset_gen_prompt = ""
                            if not is_video_prompt : # Refining an image prompt
                                ref_req_prompt_for_gemini = create_visual_regeneration_prompt(prompt_to_edit, fb_visual, scene_content_display,
                                    st.session_state.character_definitions, st.session_state.global_style_additions)
                                try:
                                    new_asset_gen_prompt = st.session_state.gemini_handler.refine_image_prompt_from_feedback(ref_req_prompt_for_gemini)
                                    st.session_state.scene_prompts[i_main] = new_asset_gen_prompt
                                    s_visual_regen.update(label="Image prompt refined by Gemini! Regenerating asset...", state="running")
                                except Exception as e_gemini_refine:
                                    s_visual_regen.update(label=f"Error refining prompt: {e_gemini_refine}", state="error");
                                    logger.error(f"Visual prompt refinement error: {e_gemini_refine}", exc_info=True)
                                    continue # Skip asset generation if prompt refinement failed
                            else: # For video, or auto choosing video, reconstruct the prompt
                                new_asset_gen_prompt = construct_text_to_video_prompt(scene_content_display, st.session_state.character_definitions, st.session_state.global_style_additions)
                                st.session_state.scene_prompts[i_main] = new_asset_gen_prompt
                                s_visual_regen.update(label="Video prompt reconstructed! Regenerating asset...", state="running")


                            v_num = 1
                            if current_asset_data and current_asset_data.get('path') and os.path.exists(current_asset_data['path']):
                                try: b,_=os.path.splitext(os.path.basename(current_asset_data['path'])); v_num = int(b.split('_v')[-1])+1 if '_v' in b else 2
                                except: v_num=2
                            else: v_num = 1

                            # <<< MODIFIED START >>> : Call new function
                            # Pass the current scene_content_display as its prompt might have changed.
                            # User asset type choice from the scene data for consistency
                            if generate_asset_for_scene_core(i_main, st.session_state.story_treatment_scenes[i_main], version=v_num, user_selected_asset_type=user_asset_type_choice):
                                s_visual_regen.update(label="Asset Updated! 🎉", state="complete", expanded=False)
                            else: s_visual_regen.update(label="Prompt updated, asset regeneration failed.", state="complete", expanded=False)
                            # <<< MODIFIED END >>>
                            st.rerun()
                    else: st.warning("Please provide feedback.")
        st.markdown("---")

    # Video Assembly Button
    # <<< MODIFIED START >>> : Check generated_scene_assets and use its data
    if st.session_state.story_treatment_scenes and any(asset_info and not asset_info.get('error') and asset_info.get('path') for asset_info in st.session_state.generated_scene_assets if asset_info 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:
                assets_for_video_assembly = []
                for i_v, sc_c in enumerate(st.session_state.story_treatment_scenes):
                    asset_info = st.session_state.generated_scene_assets[i_v] if i_v < len(st.session_state.generated_scene_assets) else None
                    if asset_info and not asset_info.get('error') and asset_info.get('path') and os.path.exists(asset_info['path']):
                        assets_for_video_assembly.append({
                            'path': asset_info['path'],
                            'type': asset_info.get('type', 'image'), # Default to image if type missing
                            'scene_num': sc_c.get('scene_number', i_v + 1),
                            'key_action': sc_c.get('key_plot_beat', ''),
                            'duration': sc_c.get('user_scene_duration_secs', DEFAULT_SCENE_DURATION_SECS) # Use user-set duration
                        })
                        status_vid.write(f"Adding Scene {sc_c.get('scene_number', i_v + 1)} ({asset_info.get('type')}).")
                    else:
                        logger.warning(f"Skipping Scene {sc_c.get('scene_number', i_v+1)} for video: No valid asset.")

                if assets_for_video_assembly:
                    status_vid.write("Calling video engine...");
                    st.session_state.video_path = st.session_state.visual_engine.assemble_animatic_from_assets( # Changed method name
                        asset_data_list=assets_for_video_assembly, # Pass the list of asset dicts
                        overall_narration_path=st.session_state.overall_narration_audio_path,
                        output_filename="cinegen_ultra_animatic.mp4",
                        fps=24
                    )
                    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 assets for video assembly.", state="error", expanded=False); logger.warning("No valid assets found for video assembly.")
    elif st.session_state.story_treatment_scenes: st.info("Generate visual assets before assembling video.")
    # <<< MODIFIED END >>>

    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")
            st.download_button(label="Download Ultra Animatic", data=video_bytes, 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")