import os
import gc
import torch
import streamlit as st
import tempfile
import json
import subprocess
import shutil
from datetime import datetime
from io import BytesIO
import random
from PIL import Image

# --- Hugging Face Model Libraries (Local Models) ---
from transformers import AutoTokenizer, AutoModelForCausalLM
from parler_tts import ParlerTTSForConditionalGeneration

# --- Google Generative AI (Gemini API) ---
try:
    from google import generativeai as genai
    from google.generativeai import types as genai_types # For GenerateContentConfig
    google_gemini_sdk_available = True
except ImportError:
    google_gemini_sdk_available = False
    # Error will be handled in UI

# --- Application Configuration ---
st.set_page_config(layout="wide", page_title="🌟 AI POV Story Weaver v2")

# --- Model IDs ---
SCRIPT_LLM_MODEL_ID = "openai-community/gpt2-medium" # Stand-in for "Tinglama"
TTS_MODEL_ID = "parler-tts/parler-tts-mini-v1.1"

# --- Gemini API Configuration (from Streamlit Secrets) ---
GEMINI_API_KEY = st.secrets.get("GEMINI_API_KEY")
GEMINI_IMAGE_MODEL_ID = st.secrets.get("GEMINI_IMAGE_MODEL_ID") # User's specified model

# --- Hugging Face Cache ---
CACHE_DIR = os.path.join(tempfile.gettempdir(), "hf_cache_story_weaver_v2")
os.makedirs(CACHE_DIR, exist_ok=True)
# (Setting environment variables for HF cache)
os.environ['HUGGINGFACE_HUB_CACHE'] = CACHE_DIR
os.environ['HF_HOME'] = CACHE_DIR
os.environ['TRANSFORMERS_CACHE'] = CACHE_DIR

# --- Session State Initialization ---
if 'run_id' not in st.session_state:
    st.session_state.run_id = datetime.now().strftime("%Y%m%d_%H%M%S")
if 'generated_data' not in st.session_state:
    st.session_state.generated_data = None
if 'temp_base_dir' not in st.session_state:
    st.session_state.temp_base_dir = None

# --- Utility Functions (largely same as before) ---
def get_session_temp_dir():
    if st.session_state.temp_base_dir and os.path.exists(st.session_state.temp_base_dir):
        return st.session_state.temp_base_dir
    base_dir = os.path.join(tempfile.gettempdir(), f"story_weaver_v2_run_{st.session_state.run_id}")
    os.makedirs(base_dir, exist_ok=True)
    st.session_state.temp_base_dir = base_dir
    return base_dir

def cleanup_temp_files():
    path_to_clean = st.session_state.get("temp_base_dir")
    if path_to_clean and os.path.exists(path_to_clean):
        try:
            shutil.rmtree(path_to_clean)
            st.session_state.temp_base_dir = None
        except Exception as e:
            st.warning(f"Warning: Could not clean up temp dir {path_to_clean}: {e}")

def clear_gpu_cache():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

# --- Model Loading Functions (Cached) ---
@st.cache_resource
def load_script_llm_resources(model_id):
    st.write(f"Loading LLM for script generation: {model_id}...")
    tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=CACHE_DIR)
    model = AutoModelForCausalLM.from_pretrained(
        model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        device_map="auto", cache_dir=CACHE_DIR
    )
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
        model.config.pad_token_id = model.config.eos_token_id
    st.write("LLM for script generation loaded.")
    return model, tokenizer

@st.cache_resource
def load_tts_resources(model_id):
    st.write(f"Loading TTS model: {model_id}...")
    model = ParlerTTSForConditionalGeneration.from_pretrained(
        model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        device_map="auto", cache_dir=CACHE_DIR
    )
    tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=CACHE_DIR)
    st.write("TTS model loaded.")
    return model, tokenizer

@st.cache_resource
def get_gemini_sdk_client(_api_key): # Argument to help Streamlit caching
    """
    Returns a Gemini SDK client instance if SDK is available and API key is provided.
    This uses the client pattern from the user's snippet.
    """
    if not google_gemini_sdk_available:
        st.error("Google Generative AI SDK (`google-generativeai`) is not installed.")
        return None
    if not _api_key:
        st.error("GEMINI_API_KEY not found in secrets.toml.")
        return None
    try:
        # According to user's snippet, they instantiate client directly with API key
        client = genai.Client(api_key=_api_key)
        st.write("Gemini SDK Client initialized.")
        return client
    except Exception as e:
        st.error(f"Error initializing Gemini SDK Client: {e}")
        return None

# --- Core Generation Functions ---

def generate_story_and_prompts(main_pov_prompt: str, num_scenes: int):
    st.info(f"Generating story and image prompts for '{main_pov_prompt}'...")
    model, tokenizer = load_script_llm_resources(SCRIPT_LLM_MODEL_ID)

    # --- Enhanced Prompt for Script LLM ---
    structured_prompt = f"""
You are an expert visual storyteller and AI prompt engineer. Your task is to generate a multi-scene story based on a user's Point-of-View (POV) prompt.
The story must be divided into exactly {num_scenes} distinct scenes.

For each scene, you must provide:
1.  "scene_number": An integer representing the scene order (e.g., 1, 2, ...).
2.  "scene_narration": A short paragraph (2-4 sentences, ~30-60 words). This narration should be from the first-person POV, be engaging, and suitable for text-to-speech. Convey emotion or atmosphere where appropriate.
3.  "image_generation_prompt": A highly descriptive and creative prompt (1-3 sentences, ~40-75 words) tailored for an advanced AI image generator like Gemini Flash. This prompt should generate a single, compelling image for the scene. Include:
    *   **Subject & Action:** Clearly define the main subject(s) and what they are doing from the POV.
    *   **Setting & Environment:** Describe the location, time of day, and key environmental details.
    *   **Visual Style & Medium:** Suggest an artistic style (e.g., "photorealistic," "cinematic with dramatic lighting," "fantasy digital art," "impressionistic oil painting," "cyberpunk anime concept art," "vintage photograph").
    *   **Camera View & Composition:** Specify camera angle if important (e.g., "first-person POV looking through a visor," "low-angle shot emphasizing scale," "close-up on a mysterious object," "wide establishing shot").
    *   **Lighting & Color:** Describe the lighting conditions (e.g., "soft morning light," "neon glow," "moonlit night," "dramatic chiaroscuro") and dominant colors or color palette.
    *   **Mood & Atmosphere:** Indicate the desired feeling (e.g., "mysterious and eerie," "hopeful and adventurous," "tense and suspenseful," "serene and peaceful").
    *   **Key Details:** Mention any specific objects, textures, or elements crucial to the scene.

User's main POV prompt: "{main_pov_prompt}"

Output the result STRICTLY as a single JSON object. The JSON object should have a key "title" (a concise title derived from the main POV prompt) and a key "scenes" which is a list of scene objects. Each scene object must contain the keys "scene_number", "scene_narration", and "image_generation_prompt".

Example of a single scene object within the "scenes" list:
{{
    "scene_number": 1,
    "scene_narration": "My metallic fingers traced the glowing hieroglyphs on the alien console. A low hum resonated through the derelict starship, promising either discovery or doom.",
    "image_generation_prompt": "First-person POV of a sleek, silver robotic hand touching intricate, glowing blue hieroglyphs on a dark, alien control panel. The background shows the dimly lit, derelict interior of a starship, with faint starlight filtering through a cracked viewport. Style: Cinematic sci-fi, photorealistic textures on the robot hand and console, mysterious and suspenseful atmosphere. Focus on the interaction between hand and console."
}}

Begin JSON output now:
```json
"""

    input_ids = tokenizer.encode(structured_prompt, return_tensors="pt").to(model.device)
    estimated_output_tokens = num_scenes * 180 + 150 # Increased estimate for richer prompts
    max_new_tokens = min(estimated_output_tokens, 1200) # Slightly increased cap

    try:
        output = model.generate(
            input_ids, max_new_tokens=max_new_tokens, do_sample=True,
            temperature=0.7, top_k=60, pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id
        )
        result_text = tokenizer.decode(output, skip_special_tokens=True)
        
        json_str_content = ""
        # Try to extract JSON block, more robustly
        if "```json" in result_text:
            json_start_index = result_text.find("```json") + len("```json")
            json_end_index = result_text.rfind("```")
            if json_start_index != -1 and json_end_index != -1 and json_end_index > json_start_index:
                json_str_content = result_text[json_start_index:json_end_index].strip()
            else: # Fallback if ending ``` is missing or malformed
                json_str_content = result_text[json_start_index:].strip()
        else: # If no ```json marker, assume the relevant part starts with {
            json_start_index = result_text.find("{")
            if json_start_index != -1:
                # Try to find matching braces, simple approach
                # This is not a perfect JSON parser but a heuristic
                open_braces = 0
                potential_json_end = -1
                for i, char in enumerate(result_text[json_start_index:]):
                    if char == '{':
                        open_braces += 1
                    elif char == '}':
                        open_braces -= 1
                        if open_braces == 0:
                            potential_json_end = json_start_index + i + 1
                            break
                if potential_json_end != -1:
                    json_str_content = result_text[json_start_index:potential_json_end]
                else: # Could not find balanced braces, take a guess
                    json_str_content = result_text[json_start_index:]


        if not json_str_content:
            st.error("LLM did not produce detectable JSON content.")
            st.text_area("LLM Full Raw Output:", result_text, height=300)
            return None

        try:
            parsed_json = json.loads(json_str_content)
        except json.JSONDecodeError as e:
            st.error(f"LLM output JSON parsing error: {e}")
            st.text_area("Attempted JSON content:", json_str_content, height=200)
            st.text_area("LLM Full Raw Output (for debugging):", result_text, height=300)
            return None

        if not isinstance(parsed_json, dict) or "scenes" not in parsed_json or not isinstance(parsed_json["scenes"], list):
            st.error("LLM output JSON structure is not as expected (missing 'scenes' list or not a dict).")
            st.json(parsed_json)
            return None
        
        if len(parsed_json["scenes"]) != num_scenes:
            st.warning(f"LLM generated {len(parsed_json['scenes'])} scenes, but {num_scenes} were requested. Adjusting...")
            parsed_json["scenes"] = parsed_json["scenes"][:num_scenes]
            while len(parsed_json["scenes"]) < num_scenes: # Pad if too few (basic)
                parsed_json["scenes"].append({
                    "scene_number": len(parsed_json["scenes"]) + 1,
                    "scene_narration": "Error: Scene data missing from LLM.",
                    "image_generation_prompt": "Error: Image prompt missing from LLM."
                })

        st.success("Story and image prompts generated successfully!")
        return parsed_json

    except Exception as e:
        st.error(f"Error during LLM story/prompt generation: {e}")
        st.text_area("LLM Full Raw Output (on exception):", result_text if 'result_text' in locals() else "N/A", height=300)
        return None
    finally:
        del model; del tokenizer; clear_gpu_cache()


def generate_images_via_gemini(story_data):
    st.info("Generating images with Gemini API...")
    sdk_client = get_gemini_sdk_client(GEMINI_API_KEY) # Use the new client getter

    if not sdk_client:
        st.error("Gemini SDK Client not initialized. Cannot generate images.")
        return None

    if not GEMINI_IMAGE_MODEL_ID:
        st.error("`GEMINI_IMAGE_MODEL_ID` is not set in secrets.toml. Cannot generate images.")
        return None
    
    st.markdown(f"**Using Gemini Model for Images:** `{GEMINI_IMAGE_MODEL_ID}`")
    st.warning(f"""
        **Note on Image Generation with `{GEMINI_IMAGE_MODEL_ID}`:**
        - This uses your specified model and API call structure.
        - Image characteristics (size, style nuances) are determined by this model.
        - The 'Seed' input from the UI is not directly used in this specific Gemini API call structure.
    """)

    images_pil = []
    for i, scene_obj in enumerate(story_data["scenes"]):
        image_prompt_text = scene_obj.get("image_generation_prompt", "A beautiful, abstract scene.") # Fallback
        scene_num = scene_obj.get("scene_number", i + 1)

        st.write(f"Requesting image for Scene {scene_num} with prompt: \"{image_prompt_text[:150]}...\"")
        try:
            # --- Using user's specified Gemini calling convention ---
            response = sdk_client.models.generate_content(
                model=GEMINI_IMAGE_MODEL_ID, # Model name passed here
                contents=[image_prompt_text], # The prompt for the image
                # As per user's snippet, config might be needed by their specific model endpoint
                generation_config=genai_types.GenerateContentConfig(
                    # response_modalities=["TEXT", "IMAGE"] # This was in user's example for GenerateContentConfig
                    # However, GenerateContentConfig does not have response_modalities.
                    # If the user's model requires this, it might be an older/internal SDK version or custom handling.
                    # For safety with public SDK, I will omit it unless specified it's for GenerateContentConfig.
                    # If it's for the top-level call, it would be different.
                    # The user's snippet has it under 'config=', implying it's for GenerateContentConfig.
                    # Let's try to include it if the types allow, otherwise this might error with public SDK.
                    # Upon checking google.generativeai.types.GenerationConfig, it does not have `response_modalities`.
                    # The user's example had `config=types.GenerateContentConfig(response_modalities=["TEXT", "IMAGE"])`
                    # This structure means `response_modalities` is an argument to `GenerateContentConfig`.
                    # If their preview SDK `types.GenerateContentConfig` accepts it, this will work.
                    # Otherwise, this line will be an error with the public SDK.
                    # For now, I will try to pass it as they specified, assuming their SDK version is different.
                    # **Update:** Based on their snippet, `response_modalities` seems to be part of `GenerateContentConfig`.
                    # However, standard `google.generativeai.types.GenerateContentConfig` doesn't list it.
                    # The `generate_content` method itself in `Model` class can take `request_options` which includes `response_mime_types`.
                    # The most robust way if `response_modalities` is not a standard config param,
                    # would be to rely on the model type to produce an image, or use `response_mime_type` if the model supports it.
                    # Given the user's code snippet, I'll include it as they had it, assuming their types.py is different.
                    # THIS IS A POTENTIAL POINT OF FAILURE IF USING STANDARD PUBLIC SDK.
                    **({"response_modalities": ["TEXT", "IMAGE"]} if hasattr(genai_types.GenerateContentConfig(), 'response_modalities') else {})
                    # The above line is a Poka-yoke to attempt to add it only if the attribute exists.
                    # A simpler approach is to just try what they gave:
                    # response_modalities = ["TEXT", "IMAGE"] # This would go into GenerateContentConfig
                    # This is very specific to their stated "working code"
                ),
                # The `config` argument in `client.models.generate_content` maps to `generation_config` for the Model service.
                # The API may also have `tool_config` and `safety_settings`.
            )
            
            generated_image = None
            response_text_parts = []

            if hasattr(response, 'parts') and response.parts:
                for part in response.parts:
                    if hasattr(part, 'text') and part.text:
                        response_text_parts.append(part.text)
                    if hasattr(part, 'mime_type') and part.mime_type and part.mime_type.startswith("image/"):
                        if hasattr(part, 'inline_data') and hasattr(part.inline_data, 'data'):
                            image_bytes = part.inline_data.data
                            generated_image = Image.open(BytesIO(image_bytes))
                            st.success(f"Image for Scene {scene_num} received from Gemini.")
                            break # Found an image
            elif hasattr(response, 'text') and not generated_image: # If no parts but has text (error or text-only response)
                 response_text_parts.append(response.text)


            if generated_image:
                images_pil.append(generated_image)
                if response_text_parts:
                    st.caption(f"Accompanying text from Gemini for Scene {scene_num}: {' '.join(response_text_parts)}")
            else:
                st.warning(f"No image data explicitly found from Gemini for Scene {scene_num}.")
                if response_text_parts:
                    st.text_area(f"Gemini Text Response (Scene {scene_num}):", value=' '.join(response_text_parts), height=100)
                else:
                    st.text(f"Raw Gemini Response (Scene {scene_num}): {response}") # Log raw response if no image
                if hasattr(response, 'prompt_feedback') and response.prompt_feedback:
                     st.warning(f"Prompt Feedback for scene {scene_num}: {response.prompt_feedback}")
                images_pil.append(None)

        except Exception as e:
            st.error(f"Error generating image for Scene {scene_num} with Gemini: {e}")
            st.error(f"Model used: {GEMINI_IMAGE_MODEL_ID}. Prompt: '{image_prompt_text[:100]}...'")
            if "API key not valid" in str(e) or "PERMISSION_DENIED" in str(e):
                st.error("Gemini API Key error. Check your key and its permissions for this model.")
            elif "Could not find model" in str(e) or "MODEL_NAME_INVALID" in str(e):
                st.error(f"Gemini Model '{GEMINI_IMAGE_MODEL_ID}' not found or invalid. Verify the model name.")
            elif "response_modalities" in str(e):
                 st.error("The `response_modalities` config might not be supported by your version of `google.generativeai.types.GenerateContentConfig` or the model endpoint. This part of the code is based on your provided 'working snippet'.")
            images_pil.append(None)
    
    if not any(images_pil):
        st.error("No images were successfully generated by Gemini.")
        return None
        
    st.success("Image generation step completed.")
    return images_pil


def generate_audio_narrations(story_data):
    st.info("Generating audio narrations with ParlerTTS...")
    tts_model, tts_tokenizer = load_tts_resources(TTS_MODEL_ID)
    audio_dir = os.path.join(get_session_temp_dir(), "audio_files")
    os.makedirs(audio_dir, exist_ok=True)
    audio_file_paths = []
    description = "A clear and engaging narrator tells a story with enthusiasm."

    for i, scene_obj in enumerate(story_data["scenes"]):
        narration_text = scene_obj.get("scene_narration", "Narration unavailable.")
        scene_num = scene_obj.get("scene_number", i + 1)
        st.write(f"Generating audio for Scene {scene_num}...")
        try:
            input_ids = tts_tokenizer(description, return_tensors="pt").input_ids.to(tts_model.device)
            prompt_input_ids = tts_tokenizer(narration_text, return_tensors="pt").input_ids.to(tts_model.device)
            generation = tts_model.generate(
                input_ids=input_ids, prompt_input_ids=prompt_input_ids,
                do_sample=True, temperature=0.7, # Slightly warmer for more expression
                repetition_penalty=1.2, guidance_scale=3.0 # Experiment with ParlerTTS params
            ).to(torch.float32)
            audio_waveform = generation.cpu().numpy().squeeze()
            file_path = os.path.join(audio_dir, f"s_{scene_num}_audio.wav")
            sf.write(file_path, audio_waveform, tts_model.config.sampling_rate)
            audio_file_paths.append(file_path)
            st.success(f"Audio for Scene {scene_num} created.")
        except Exception as e:
            st.error(f"Audio error (Scene {scene_num}): {e}")
            audio_file_paths.append(None)
    del tts_model; del tts_tokenizer; clear_gpu_cache()
    st.success("Audio narration step completed.")
    return audio_file_paths


def create_final_video(image_pil_objects, audio_paths):
    # (This function remains largely the same as v1, ensure paths and checks are robust)
    st.info("Creating final video...")
    if not image_pil_objects or not audio_paths or len(image_pil_objects) != len(audio_paths):
        st.error("Asset mismatch for video. Cannot create."); return None
    try: subprocess.run(["ffmpeg", "-version"], capture_output=True, check=True)
    except: st.error("FFMPEG not found."); return None

    temp_image_dir = os.path.join(get_session_temp_dir(), "vid_frames")
    os.makedirs(temp_image_dir, exist_ok=True)
    img_paths_for_vid = []
    for idx, img_pil in enumerate(image_pil_objects):
        if img_pil:
            p = os.path.join(temp_image_dir, f"f_{idx:03d}.png"); img_pil.save(p)
            img_paths_for_vid.append(p)
        else: img_paths_for_vid.append(None)

    temp_clips_dir = os.path.join(get_session_temp_dir(), "temp_vid_clips")
    os.makedirs(temp_clips_dir, exist_ok=True)
    vid_clip_paths, valid_clips = [], 0

    for i, (img_p, aud_p) in enumerate(zip(img_paths_for_vid, audio_paths)):
        s_num = i + 1
        if not (img_p and aud_p): st.warning(f"Skipping Scene {s_num} in video (missing asset)."); continue
        try:
            aud_info = sf.info(aud_p); aud_dur = aud_info.duration
            if aud_dur < 0.5: aud_dur = 0.5 # Min clip duration
            clip_out_p = os.path.join(temp_clips_dir, f"c_{s_num:03d}.mp4")
            cmd = [
                "ffmpeg", "-y", "-loop", "1", "-i", img_p, "-i", aud_p,
                "-c:v", "libx264", "-preset", "fast", "-tune", "stillimage", # Faster preset
                "-vf", "scale=1280:720:force_original_aspect_ratio=decrease,pad=1280:720:(ow-iw)/2:(oh-ih)/2,setsar=1", # Scale and pad to 720p
                "-c:a", "aac", "-b:a", "192k", "-pix_fmt", "yuv420p",
                "-t", str(aud_dur), "-shortest", clip_out_p
            ]
            res = subprocess.run(cmd, capture_output=True, text=True)
            if res.returncode != 0: st.error(f"FFMPEG clip error (S{s_num}):\n{res.stderr}"); continue
            vid_clip_paths.append(clip_out_p); valid_clips +=1
            st.write(f"Video clip for Scene {s_num} processed.")
        except Exception as e: st.error(f"Video processing error (S{s_num}): {e}")
    
    if not vid_clip_paths or valid_clips == 0: st.error("No valid video clips. Cannot create final video."); return None

    concat_list_f = os.path.join(temp_clips_dir, "concat_list.txt")
    with open(concat_list_f, "w") as f:
        for clip_p in vid_clip_paths: f.write(f"file '{os.path.basename(clip_p)}'\n")
    
    final_vid_out_p = os.path.join(get_session_temp_dir(), "final_story_video_720p.mp4")
    concat_cmd = ["ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", concat_list_f, "-c", "copy", final_vid_out_p]
    st.write("Concatenating video clips...")
    res = subprocess.run(concat_cmd, capture_output=True, text=True, cwd=temp_clips_dir)
    if res.returncode != 0: st.error(f"FFMPEG concat error:\n{res.stderr}"); return None
    
    st.success("Final video created successfully!"); return final_vid_out_p

# --- Streamlit UI (largely same, check key names and logic) ---
st.title("🌟 AI POV Story Weaver v2 ✨")
st.markdown("Craft unique POV stories with refined AI-generated scripts, Gemini images, voiceovers, and a final video!")
st.markdown("---")

with st.sidebar:
    st.header("🛠️ Story Configuration")
    user_main_prompt = st.text_area(
        "Enter your main POV story idea:",
        st.session_state.get("user_main_prompt_val", "POV: I'm a lone astronaut discovering an ancient, bioluminescent forest on a new planet."),
        height=120, key="main_prompt_input_v2"
    )
    num_scenes_input = st.slider(
        "Number of Scenes:", 1, 4, # Max 4 for resource management
        st.session_state.get("num_scenes_input_val", 2), key="num_scenes_slider_v2"
    )
    st.caption(f"Script LLM: `{SCRIPT_LLM_MODEL_ID}`")
    st.caption(f"TTS Model: `{TTS_MODEL_ID}`")
    st.caption(f"Image Model (Gemini): `{GEMINI_IMAGE_MODEL_ID}` (from secrets)")

    if not google_gemini_sdk_available: st.error("Google SDK missing (`pip install google-generativeai`)")
    if not GEMINI_API_KEY: st.error("`GEMINI_API_KEY` not set in secrets.")
    if not GEMINI_IMAGE_MODEL_ID or "your-gemini" in GEMINI_IMAGE_MODEL_ID or "flash-preview" not in GEMINI_IMAGE_MODEL_ID :
         st.warning(f"Model ID '{GEMINI_IMAGE_MODEL_ID}' might be placeholder or not your specific preview model. Ensure it's correct in secrets for image generation.")

    st.markdown("---")
    can_generate = google_gemini_sdk_available and GEMINI_API_KEY and GEMINI_IMAGE_MODEL_ID
    if st.button("🎬 Weave My Story! (v2)", type="primary", use_container_width=True, disabled=not can_generate):
        st.session_state.run_id = datetime.now().strftime("%Y%m%d_%H%M%S")
        cleanup_temp_files()
        st.session_state.generated_data = {}
        st.session_state.user_main_prompt_val = user_main_prompt
        st.session_state.num_scenes_input_val = num_scenes_input
        st.session_state.trigger_generation_v2 = True
    
    st.markdown("---")
    if st.button("🧹 Clear All & Reset", use_container_width=True):
        cleanup_temp_files()
        keys_to_clear = ['generated_data', 'trigger_generation_v2', 'user_main_prompt_val', 'num_scenes_input_val']
        for key in keys_to_clear:
            if key in st.session_state: del st.session_state[key]
        st.cache_resource.clear()
        st.success("Cleared temp files, state, and model cache. Reload on next run."); st.rerun()

# --- Main Area for Generation and Display ---
if st.session_state.get("trigger_generation_v2"):
    with st.spinner("📜 Phase 1: Generating enhanced story script and image prompts..."):
        story_json_data = generate_story_and_prompts(
            st.session_state.user_main_prompt_val, st.session_state.num_scenes_input_val
        )
        st.session_state.generated_data['story_json'] = story_json_data

    if story_json_data:
        # (Display logic for story, images, audio, video - similar to previous, ensuring keys match)
        st.header("📜 Generated Story & Image Prompts")
        st.json(story_json_data)
        # ... (Download button for story_json) ...
        st.markdown("---")

        with st.spinner(f"🎨 Phase 2: Generating images with Gemini ({GEMINI_IMAGE_MODEL_ID})..."):
            pil_images = generate_images_via_gemini(story_json_data)
            st.session_state.generated_data['pil_images'] = pil_images
        
        if pil_images and any(pil_images): # Check if list is not empty AND contains at least one image
            st.header("🖼️ Generated Images")
            # ... (Display and download buttons for pil_images) ...
            cols = st.columns(min(len(pil_images), 3)) 
            for i, img in enumerate(pil_images):
                if img:
                    with cols[i % len(cols)]:
                        st.image(img, caption=f"Scene {story_json_data['scenes'][i].get('scene_number', i+1)}")
                        # ... (download button for img) ...
            st.markdown("---")

            with st.spinner("🔊 Phase 3: Generating audio narrations..."):
                audio_paths = generate_audio_narrations(story_json_data)
                st.session_state.generated_data['audio_paths'] = audio_paths
            
            if audio_paths and any(audio_paths):
                st.header("🎤 Generated Audio Narrations")
                # ... (Display and download buttons for audio_paths) ...
                st.markdown("---")

                if st.session_state.generated_data.get('pil_images') and st.session_state.generated_data.get('audio_paths'):
                    with st.spinner("🎬 Final Phase: Weaving the video masterpiece..."):
                        final_video_path = create_final_video(
                            st.session_state.generated_data['pil_images'],
                            st.session_state.generated_data['audio_paths']
                        )
                        st.session_state.generated_data['final_video_path'] = final_video_path
                    
                    if final_video_path:
                        st.header("🎉 Your Story Video is Ready! 🎉")
                        st.video(final_video_path)
                        # ... (Download button for final_video_path) ...
                        st.balloons()
                    else: st.error("Video creation failed. Check FFMPEG logs if any were shown.")
                else: st.warning("Skipping video: not all images or audio were generated.")
            else: st.error("Audio generation failed. Cannot proceed to video.")
        else: st.error("Image generation failed (no images returned). Cannot proceed.")
    else: st.error("Story script generation failed. Cannot proceed.")
    st.session_state.trigger_generation_v2 = False # Reset trigger

elif not st.session_state.get("user_main_prompt_val"):
    st.info("👋 Welcome to the AI Story Weaver v2! Configure your story in the sidebar and click 'Weave My Story!'")