app.py

# Copyright 2025 Google LLC. Based on work by Yousif Ahmed.
# Concept: ChronoWeave - Branching Narrative Generation
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0

import streamlit as st
import google.generativeai as genai
import os
import json
import numpy as np
from io import BytesIO
import time
import wave
import contextlib
import asyncio
import uuid # For unique filenames
import shutil # For cleaning up temp dirs

# Image handling
from PIL import Image

# Video and audio processing
from moviepy.editor import ImageClip, AudioFileClip, CompositeVideoClip, concatenate_videoclips

# Type hints
import typing_extensions as typing

# Async support for Streamlit/Google API
import nest_asyncio
nest_asyncio.apply() # Apply patch for asyncio in environments like Streamlit/Jupyter

# --- Configuration ---
st.set_page_config(page_title="ChronoWeave", layout="wide")
st.title("🌀 ChronoWeave: Branching Narrative Generator")
st.markdown("""
Generate multiple, branching story timelines from a single theme using AI.
Based on the work of Yousif Ahmed. Copyright 2025 Google LLC.
""")

# --- Constants ---
MODEL = "models/gemini-1.5-flash" # Or other suitable text model supporting JSON
# Using v1alpha for the Live API for audio output.
AUDIO_MODEL_VERSION = 'v1alpha' # Must be alpha for audio modality
IMAGE_MODEL_ID = "imagen-3" # Or your preferred Imagen model "imagen-3.0-generate-002"

# --- API Key Handling ---
try:
    # Preferred way to handle secrets in Streamlit sharing/HF Spaces
    GOOGLE_API_KEY = st.secrets["GOOGLE_API_KEY"]
    os.environ['GOOGLE_API_KEY'] = GOOGLE_API_KEY
except KeyError:
    st.error("🚨 Google API Key not found! Please add it as a Secret named 'GOOGLE_API_KEY' in your Hugging Face Space settings.", icon="🚨")
    st.stop() # Halt execution if no key

# --- Initialize Google Client ---
try:
    # Initialize the client with the API key
    genai.configure(api_key=GOOGLE_API_KEY)

    # Create separate clients or configure one for different API versions if needed
    # Client for Text/Imagen (standard API)
    client_standard = genai.GenerativeModel(MODEL)
    # Client for Live Audio (v1alpha) - requires different client init
    client_live = genai.Client(
         client_options={'api_endpoint': f'{AUDIO_MODEL_VERSION}.generativelanguage.googleapis.com'}
    )
    # Note: As of recent updates, genai.configure might handle this better,
    # but separating clients or explicitly setting endpoints can be more robust.
    # Adjust based on the library version and observed behavior.


except Exception as e:
    st.error(f"🚨 Failed to initialize Google AI Client: {e}", icon="🚨")
    st.stop()


# --- Define Structured Output Schemas ---
class StorySegment(typing.TypedDict):
    scene_id: int
    image_prompt: str
    audio_text: str
    character_description: str
    timeline_visual_modifier: typing.Optional[str]

class Timeline(typing.TypedDict):
    timeline_id: int
    divergence_reason: str
    segments: list[StorySegment]

class ChronoWeaveResponse(typing.TypedDict):
    core_theme: str
    timelines: list[Timeline]
    total_scenes_per_timeline: int

# --- Helper Functions ---

@contextlib.contextmanager
def wave_file(filename, channels=1, rate=24000, sample_width=2):
    """Context manager to write WAV files."""
    with wave.open(filename, "wb") as wf:
        wf.setnchannels(channels)
        wf.setsampwidth(sample_width)
        wf.setframerate(rate)
        yield wf

async def generate_audio_live_async(api_text, output_filename):
    """Generates audio using Gemini Live API (async version)."""
    collected_audio = bytearray()
    st.write(f"🎙️ Generating audio for: '{api_text[:50]}...'") # Log start

    try:
        # Use the 'client_live' specifically configured for v1alpha
        live_model = client_live.get_model(f"models/gemini-1.5-flash") # Specify model within the live client context

        config = {
            "response_modalities": ["AUDIO"]
        }
        # Connect to the Live API using the live client.
        async with live_model.connect(config=config) as session:
            await session.send_request([api_text]) # Simpler send for single prompt
            async for response in session.stream_content():
                 if response.audio_chunk:
                    collected_audio.extend(response.audio_chunk.data)

        if not collected_audio:
             st.warning(f"⚠️ No audio data received for: '{api_text[:50]}...'")
             return None # Indicate failure

        audio_bytes = bytes(collected_audio)
        # Write the collected audio bytes into a WAV file.
        with wave_file(output_filename) as wf:
            wf.writeframes(audio_bytes)
        st.write(f"   ✅ Audio saved: {os.path.basename(output_filename)}")
        return output_filename
    except Exception as e:
        st.error(f"   ❌ Audio generation failed for '{api_text[:50]}...': {e}", icon="🚨")
        return None


def generate_story_sequence_chrono(theme: str, num_scenes: int, num_timelines: int, divergence_prompt: str = "") -> ChronoWeaveResponse | None:
    """Generates branching story sequences using Gemini structured output."""
    st.write(f"📚 Generating {num_timelines} timeline(s) for theme: '{theme}'...")
    divergence_instruction = f"Introduce divergence between timelines. {divergence_prompt}" if divergence_prompt else "Introduce natural points of divergence between timelines after the first scene or two."

    prompt = f'''
    As an expert narrative designer, create a branching story based on the theme: "{theme}".
    Generate exactly {num_timelines} distinct timelines, each containing exactly {num_scenes} scenes.
    Each scene should be approximately 5-10 seconds long when narrated.

    {divergence_instruction} Clearly state the reason for divergence for each timeline after the first.

    For each scene in each timeline, provide:
    - scene_id: An integer starting from 0 for the scene number within its timeline.
    - image_prompt: A concise (15-25 words) description for an image generation model. Focus on visual details, characters (animals/objects only, NO PEOPLE), background, and action. Maintain a consistent 'kids animation style' (e.g., simple, rounded shapes, bright colors) across all scenes and timelines unless specified by a timeline_visual_modifier.
    - audio_text: A single, engaging sentence of narration or dialogue for the scene (max 25 words).
    - character_description: Brief description of recurring characters (names, key features) mentioned in *this specific scene's image prompt*. Keep consistent within a timeline. (Max 30 words).
    - timeline_visual_modifier: (Optional, string or null) A *brief* hint if this timeline should have a slightly different visual feel from this scene onwards (e.g., "slightly darker lighting", "more cluttered background", "character looks worried"). Keep it subtle. Use null if no specific modifier.

    Constraint: Ensure the output strictly adheres to the following JSON schema. Do not include preamble or explanations outside the JSON structure. Respond ONLY with the JSON object.

    JSON Schema:
    {{
      "type": "object",
      "properties": {{
        "core_theme": {{"type": "string"}},
        "timelines": {{
          "type": "array",
          "items": {{
            "type": "object",
            "properties": {{
              "timeline_id": {{"type": "integer"}},
              "divergence_reason": {{"type": "string"}},
              "segments": {{
                "type": "array",
                "items": {{
                  "type": "object",
                  "properties": {{
                    "scene_id": {{"type": "integer"}},
                    "image_prompt": {{"type": "string"}},
                    "audio_text": {{"type": "string"}},
                    "character_description": {{"type": "string"}},
                    "timeline_visual_modifier": {{"type": ["string", "null"]}}
                  }},
                  "required": ["scene_id", "image_prompt", "audio_text", "character_description", "timeline_visual_modifier"]
                }}
              }}
            }},
            "required": ["timeline_id", "divergence_reason", "segments"]
          }}
        }},
        "total_scenes_per_timeline": {{"type": "integer"}}
      }},
      "required": ["core_theme", "timelines", "total_scenes_per_timeline"]
    }}
    '''

    try:
        response = client_standard.generate_content(
            contents=prompt,
            generation_config=genai.types.GenerationConfig(
                 response_mime_type="application/json",
                 # Optional: Add temperature, etc. if needed
            )
            # The schema can also be passed via generation_config in some versions/models
            # config={
            #     'response_mime_type': 'application/json',
            #     'response_schema': ChronoWeaveResponse # Pass the TypedDict directly
            # }
        )

        # Debugging: Print raw response text
        # st.text_area("Raw Gemini Response:", response.text, height=200)

        story_data = json.loads(response.text) # Assuming response.text contains the JSON string
        st.success("✅ Story structure generated successfully!")
        # Basic validation (can be more thorough)
        if 'timelines' in story_data and isinstance(story_data['timelines'], list):
             # Further validation could check segment structure, etc.
             return story_data # Return the parsed dictionary
        else:
             st.error("🚨 Generated story data is missing the 'timelines' list.", icon="🚨")
             return None

    except json.JSONDecodeError as e:
        st.error(f"🚨 Failed to decode JSON response from Gemini: {e}", icon="🚨")
        st.text_area("Problematic Response Text:", response.text if 'response' in locals() else "No response object.", height=150)
        return None
    except Exception as e:
        st.error(f"🚨 Error generating story sequence: {e}", icon="🚨")
        # Log the prompt potentially? Be careful with sensitive data if applicable.
        # st.text_area("Failed Prompt:", prompt, height=200)
        return None


def generate_image_imagen(prompt: str, aspect_ratio: str = "1:1") -> Image.Image | None:
    """Generates an image using Imagen."""
    st.write(f"🖼️ Generating image for: '{prompt[:60]}...'")
    try:
        # Use the standard client's dedicated image generation method
        response = client_standard.generate_content(
            f"Generate an image with the following prompt, ensuring a child-friendly animation style and NO human figures: {prompt}",
             generation_config=genai.types.GenerationConfig(
                 candidate_count=1, # Generate one image
                 # Imagen specific parameters are often passed differently or rely on model defaults
                 # Check documentation for precise Imagen control via the unified API
             ),
             # If the model/API version requires specific image parameters:
             # tools=[genai.ImageParams(model=IMAGE_MODEL_ID, number_of_images=1, aspect_ratio=aspect_ratio, person_generation="DONT_ALLOW")]
        )

        # Accessing image data might vary slightly depending on API response structure
        # This assumes response.parts contains the image data if successful
        if response.parts and response.parts[0].inline_data:
            image_bytes = response.parts[0].inline_data.data
            image = Image.open(BytesIO(image_bytes))
            st.write("   ✅ Image generated.")
            return image
        else:
            # Check for safety blocks or other reasons for failure
            if response.prompt_feedback.block_reason:
                 st.warning(f"   ⚠️ Image generation blocked for prompt '{prompt[:60]}...'. Reason: {response.prompt_feedback.block_reason}", icon="⚠️")
            else:
                 st.warning(f"   ⚠️ No image data received for prompt '{prompt[:60]}...'.", icon="⚠️")
            # Debugging: st.write(response)
            return None

    except Exception as e:
        st.error(f"   ❌ Image generation failed for '{prompt[:60]}...': {e}", icon="🚨")
        return None


# --- Streamlit UI Elements ---
st.sidebar.header("Configuration")

# API Key display/check (already handled above, but sidebar is a good place)
if GOOGLE_API_KEY:
    st.sidebar.success("Google API Key Loaded!", icon="✅")
else:
    st.sidebar.error("Google API Key Missing!", icon="🚨")

theme = st.sidebar.text_input("Story Theme:", "A curious squirrel finds a shiny object")
num_scenes = st.sidebar.slider("Scenes per Timeline:", min_value=2, max_value=7, value=3)
num_timelines = st.sidebar.slider("Number of Timelines:", min_value=1, max_value=4, value=2)
divergence_prompt = st.sidebar.text_input("Divergence Hint (Optional):", placeholder="e.g., What if it started raining?")
aspect_ratio = st.sidebar.selectbox("Image Aspect Ratio:", ["1:1", "16:9", "9:16"], index=0)

generate_button = st.sidebar.button("✨ Generate ChronoWeave ✨", type="primary", disabled=(not GOOGLE_API_KEY))

st.sidebar.markdown("---")
st.sidebar.info("Note: Generation can take several minutes depending on settings.")

# --- Main Logic ---
if generate_button:
    if not theme:
        st.error("Please enter a story theme.", icon="👈")
    else:
        # Create a unique temporary directory for this run
        run_id = str(uuid.uuid4())
        temp_dir = os.path.join(".", f"chrono_temp_{run_id}") # Create in current dir
        os.makedirs(temp_dir, exist_ok=True)
        st.write(f"Working directory: {temp_dir}")

        final_video_paths = {} # To store {timeline_id: video_path}

        with st.spinner("Generating narrative structure..."):
            chrono_data = generate_story_sequence_chrono(theme, num_scenes, num_timelines, divergence_prompt)

        if chrono_data and 'timelines' in chrono_data:
            st.success(f"Found {len(chrono_data['timelines'])} timelines. Processing each...")

            all_timelines_successful = True # Flag to track if all timelines worked

            # Use st.status for detailed progress
            with st.status("Generating assets and composing videos...", expanded=True) as status:

                for timeline in chrono_data['timelines']:
                    timeline_id = timeline['timeline_id']
                    divergence = timeline['divergence_reason']
                    segments = timeline['segments']
                    st.subheader(f"Timeline {timeline_id}: {divergence}")

                    temp_image_files = []
                    temp_audio_files = []
                    video_clips = []
                    timeline_successful = True # Flag for this specific timeline

                    for i, segment in enumerate(segments):
                        status.update(label=f"Processing Timeline {timeline_id}, Scene {i+1}/{num_scenes}...")
                        scene_id = segment['scene_id']
                        image_prompt = segment['image_prompt']
                        audio_text = segment['audio_text']
                        char_desc = segment['character_description']
                        vis_mod = segment['timeline_visual_modifier']

                        st.write(f"--- Scene {i+1} (T{timeline_id}) ---")
                        st.write(f"*   **Image Prompt:** {image_prompt}" + (f" (Modifier: {vis_mod})" if vis_mod else ""))
                        st.write(f"*   **Audio Text:** {audio_text}")
                        # st.write(f"*   Character Desc: {char_desc}") # Can be verbose

                        # --- Image Generation ---
                        combined_prompt = f"{image_prompt} {char_desc}"
                        if vis_mod:
                            combined_prompt += f" Style hint: {vis_mod}"

                        generated_image = generate_image_imagen(combined_prompt, aspect_ratio)

                        if generated_image:
                            image_path = os.path.join(temp_dir, f"t{timeline_id}_s{i}_image.png")
                            generated_image.save(image_path)
                            temp_image_files.append(image_path)
                            st.image(generated_image, width=200) # Show thumbnail
                        else:
                            st.warning(f"Skipping scene {i+1} in timeline {timeline_id} due to image generation failure.")
                            timeline_successful = False
                            continue # Skip to next segment if image fails

                        # --- Audio Generation ---
                        # Add negative prompt to prevent conversational filler
                        audio_negative_prompt = "Narrate the following sentence directly, with expression, without any introduction or closing remarks like 'Okay' or 'Here is the narration'. Just read the sentence:"
                        full_audio_prompt = f"{audio_negative_prompt}\n{audio_text}"
                        audio_path = os.path.join(temp_dir, f"t{timeline_id}_s{i}_audio.wav")

                        # Run the async audio generation function
                        try:
                           generated_audio_path = asyncio.run(generate_audio_live_async(full_audio_prompt, audio_path))
                        except Exception as e:
                            st.error(f"Asyncio error during audio gen: {e}")
                            generated_audio_path = None


                        if generated_audio_path:
                            temp_audio_files.append(generated_audio_path)
                            # st.audio(generated_audio_path) # Optional: Preview audio
                        else:
                            st.warning(f"Skipping video clip for scene {i+1} in timeline {timeline_id} due to audio generation failure.")
                            # Clean up the image file for this failed scene segment
                            if os.path.exists(image_path):
                                os.remove(image_path)
                                temp_image_files.remove(image_path)
                            timeline_successful = False
                            continue # Skip making video clip if audio fails


                        # --- Create Video Clip ---
                        try:
                            st.write("   🎬 Creating video clip...")
                            audio_clip = AudioFileClip(generated_audio_path)
                            # Ensure PIL Image is used if needed, or numpy array directly
                            np_image = np.array(Image.open(image_path))
                            # Create ImageClip, ensure duration matches audio
                            image_clip = ImageClip(np_image).set_duration(audio_clip.duration)

                            # Handle potential size mismatch if needed (resize image_clip or set size explicitly)
                            # image_clip = image_clip.resize(width=...)

                            composite_clip = image_clip.set_audio(audio_clip) # Simpler composition
                            video_clips.append(composite_clip)
                            st.write("      ✅ Clip created.")
                        except Exception as e:
                             st.error(f"      ❌ Failed to create video clip for scene {i+1} (T{timeline_id}): {e}", icon="🚨")
                             timeline_successful = False
                             # Don't break the whole timeline, just skip this clip maybe? Or mark timeline as failed.


                    # --- Assemble Timeline Video ---
                    if video_clips and timeline_successful: # Only assemble if clips were made and no major errors
                        status.update(label=f"Composing final video for Timeline {timeline_id}...")
                        st.write(f"🎞️ Assembling final video for Timeline {timeline_id}...")
                        try:
                            final_timeline_video = concatenate_videoclips(video_clips, method="compose")
                            output_filename = os.path.join(temp_dir, f"timeline_{timeline_id}_final_video.mp4")
                            # Use 'libx264' for broader compatibility, specify audio codec
                            final_timeline_video.write_videofile(output_filename, fps=24, codec='libx264', audio_codec='aac')
                            final_video_paths[timeline_id] = output_filename
                            st.success(f"   ✅ Video for Timeline {timeline_id} saved: {os.path.basename(output_filename)}")

                            # Close clips to release resources
                            for clip in video_clips:
                                if hasattr(clip, 'close'): clip.close()
                                if hasattr(clip, 'audio') and hasattr(clip.audio, 'close'): clip.audio.close()
                            if hasattr(final_timeline_video, 'close'): final_timeline_video.close()


                        except Exception as e:
                            st.error(f"   ❌ Failed to write final video for Timeline {timeline_id}: {e}", icon="🚨")
                            all_timelines_successful = False
                    elif not video_clips:
                         st.warning(f"No video clips were successfully generated for Timeline {timeline_id}. Skipping final video assembly.")
                         all_timelines_successful = False
                    else:
                         st.warning(f"Timeline {timeline_id} encountered errors. Skipping final video assembly.")
                         all_timelines_successful = False


                    # Intermediate cleanup for the timeline (optional, helps manage files)
                    # for file in temp_audio_files:
                    #     if os.path.exists(file): os.remove(file)
                    # for file in temp_image_files:
                    #     if os.path.exists(file): os.remove(file)

                # Final status update
                if all_timelines_successful and final_video_paths:
                    status.update(label="ChronoWeave Generation Complete!", state="complete", expanded=False)
                elif final_video_paths:
                     status.update(label="ChronoWeave Generation Partially Complete (some errors occurred).", state="warning", expanded=False)
                else:
                    status.update(label="ChronoWeave Generation Failed.", state="error", expanded=False)


            # --- Display Results ---
            st.header("Generated Timelines")
            if final_video_paths:
                sorted_timeline_ids = sorted(final_video_paths.keys())
                for timeline_id in sorted_timeline_ids:
                    video_path = final_video_paths[timeline_id]
                    # Find matching timeline divergence reason
                    reason = "Unknown"
                    for t in chrono_data.get('timelines', []):
                         if t.get('timeline_id') == timeline_id:
                             reason = t.get('divergence_reason', 'N/A')
                             break
                    st.subheader(f"Timeline {timeline_id}: {reason}")
                    try:
                        video_file = open(video_path, 'rb')
                        video_bytes = video_file.read()
                        st.video(video_bytes)
                        video_file.close()
                    except FileNotFoundError:
                         st.error(f"Could not find video file: {video_path}", icon="🚨")
                    except Exception as e:
                         st.error(f"Could not display video {video_path}: {e}", icon="🚨")
            else:
                st.warning("No final videos were successfully generated.")

            # --- Cleanup ---
            st.write("Cleaning up temporary files...")
            try:
                shutil.rmtree(temp_dir)
                st.write("   ✅ Temporary files removed.")
            except Exception as e:
                st.warning(f"   ⚠️ Could not remove temporary directory {temp_dir}: {e}", icon="⚠️")


        elif not chrono_data:
            st.error("Story generation failed. Cannot proceed.", icon="🛑")
        else:
            # This case might happen if chrono_data is returned but is malformed (e.g., no 'timelines' key)
            st.error("Story data seems malformed. Cannot proceed.", icon="🛑")
            # st.json(chrono_data) # Display the problematic data

else:
    st.info("Configure settings in the sidebar and click 'Generate ChronoWeave'")