#!/usr/bin/env python3
"""
Stage 7: Grid Generation

Two-step process:
1. Call Claude to create a NanoBanana prompt (nanobanana_prompt.txt) + generation_config.json
2. Call NanoBanana API to generate the grid image

Creates 07-grids/XXX/nanobanana_prompt.txt + generation_config.json + grid.png
"""

import subprocess
import sys
import json
from pathlib import Path
from datetime import datetime, timezone

# Add parent to path for lib imports
sys.path.insert(0, str(Path(__file__).parent.parent))

# Add repo scripts dir for nanobanana import
repo_root = Path(__file__).parent.parent.parent.parent
sys.path.insert(0, str(repo_root / "scripts"))

from lib import (
    save_metadata,
    ensure_dir,
    setup_generation_log,
    log_event,
    claude_env,
    build_reference_record,
    scale_down_references
)
from lib.path_utils_v3 import get_next_iteration
from lib.feedback_history_v5 import get_feedback_history_for_claude
from lib.reference_analysis import build_analyze_prefix
import reference_loader
from reference_loader import (
    is_story_mode, get_episode_characters, select_character_refs,
    get_episode_locations, get_episode_objects, select_asset_refs,
    select_asset_refs_with_fallback,
    NANOBANANA_MAX_IMAGES,
)


def execute(take_path: Path, story_name: str, inputs: dict,
            feedback: str = None, refined_from: str = None, chat_mode: bool = False,
            variation_index: int = None, total_variations: int = None,
            scale_refs: bool = True, char_ref_only: bool = False) -> dict:
    """
    Execute Stage 7: Grid Generation.

    Args:
        take_path: Path to take directory
        story_name: Story name (slug)
        inputs: Parent iterations, e.g. {"03-script": "006", "04-assets": "004", "06-characteristic": "006"}
        feedback: User feedback text (if any)
        refined_from: Iteration being refined (if refinement)
        chat_mode: Use NanoBanana chat continuation
        variation_index: Index of this variation (0-based) when generating multiple grids
        total_variations: Total number of variations being generated
    """
    logger = setup_generation_log(take_path)
    logger.section("Stage 7: Grid Generation")

    # Resolve inputs
    script_iter = inputs.get("03-script")
    assets_iter = inputs.get("04-assets")
    char_iter = inputs.get("06-characteristic")

    if not script_iter:
        return {"status": "error", "error": "No 03-script iteration in inputs"}
    if not assets_iter:
        return {"status": "error", "error": "No 04-assets iteration in inputs"}
    if not char_iter:
        return {"status": "error", "error": "No 06-characteristic iteration in inputs"}

    script_file = take_path / "03-script" / script_iter / "script.json"
    assets_file = take_path / "04-assets" / assets_iter / "assets.json"

    # Find characteristic shot (png or jpg)
    char_dir = take_path / "06-characteristic" / char_iter
    char_shot = None
    for ext in ["jpg", "png"]:
        candidate = char_dir / f"characteristic_shot.{ext}"
        if candidate.exists():
            char_shot = candidate
            break

    if not script_file.exists():
        return {"status": "error", "error": f"Script file not found: {script_file}"}
    if not assets_file.exists():
        return {"status": "error", "error": f"Assets file not found: {assets_file}"}
    if not char_shot:
        pass  # Will be caught below

    # Load video_mode from script metadata (propagated from understanding → script)
    script_metadata_file = take_path / "03-script" / script_iter / "metadata.json"
    video_mode = "standard"
    if script_metadata_file.exists():
        try:
            script_meta = json.loads(script_metadata_file.read_text())
            video_mode = script_meta.get("video_mode", "standard")
        except (json.JSONDecodeError, OSError):
            pass
    logger.info(f"Video mode: {video_mode}")

    # Determine grid layout from scene count
    try:
        script_data = json.loads(script_file.read_text())
        num_scenes = len(script_data.get("scenes", []))
    except (json.JSONDecodeError, OSError):
        num_scenes = 9
    if num_scenes > 9:
        grid_layout = "4x4"
    elif num_scenes > 4:
        grid_layout = "3x3"
    else:
        grid_layout = "2x2"
    logger.info(f"Grid layout: {grid_layout} ({num_scenes} scenes)")

    if not char_shot:
        return {"status": "error", "error": f"Characteristic shot not found in {char_dir}"}

    logger.info(f"Script: {script_file}")
    logger.info(f"Assets: {assets_file}")
    logger.info(f"Characteristic shot: {char_shot}")

    # Determine iteration number
    grid_stage = take_path / "07-grids"
    ensure_dir(grid_stage)
    iteration = get_next_iteration(grid_stage)

    iter_path = grid_stage / iteration
    ensure_dir(iter_path)

    logger.info(f"Creating iteration {iteration}")
    if refined_from:
        logger.info(f"Refined from: {refined_from}")

    # Save feedback if provided
    if feedback:
        (iter_path / "feedback.txt").write_text(feedback)
        logger.info("Saved feedback.txt")

    # Load session ID
    session_file = take_path.parent / "session.txt"
    if not session_file.exists():
        return {"status": "error", "error": "Session file not found"}

    session_id = session_file.read_text().strip()
    logger.info(f"Session ID: {session_id}")

    # ===== Determine Generation Mode =====
    # Three modes:
    # 1. "new" - Fresh generation from characteristic shot + optional feedback
    # 2. "refine" - Regenerate prompt with ALL accumulated feedback from iterations with same char shot
    # 3. "chat" - Direct chat continuation with NanoBanana (skip Claude)

    generation_mode = None

    if chat_mode and refined_from and feedback:
        # Mode 3: Chat continuation
        generation_mode = "chat"
        logger.info(f"Mode: CHAT (continuing conversation from {refined_from})")
    elif refined_from and feedback:
        # Mode 2: Refine with accumulated feedback
        generation_mode = "refine"
        logger.info(f"Mode: REFINE (collecting all feedback from {refined_from})")
    else:
        # Mode 1: Fresh generation
        generation_mode = "new"
        logger.info(f"Mode: NEW (fresh generation from characteristic shot)")

    # Collect feedback for refine mode
    if generation_mode == "refine":
        # Get ALL feedback from previous iterations with the same characteristic shot
        all_feedback = []
        grid_stage_dir = take_path / "07-grids"
        if grid_stage_dir.exists():
            for iter_dir in sorted(grid_stage_dir.iterdir()):
                if not iter_dir.is_dir():
                    continue
                # Check if this iteration used the same characteristic shot
                iter_metadata = iter_dir / "metadata.json"
                if iter_metadata.exists():
                    meta = json.loads(iter_metadata.read_text())
                    if meta.get("based_on", {}).get("06-characteristic") == char_iter:
                        # Same characteristic shot - collect feedback
                        fb_file = iter_dir / "feedback.txt"
                        if fb_file.exists():
                            all_feedback.append({
                                "iteration": iter_dir.name,
                                "feedback": fb_file.read_text().strip()
                            })
        logger.info(f"Collected {len(all_feedback)} feedback entries from previous iterations")
    else:
        all_feedback = None

    story_path = take_path.parent

    # ===== MODE 3: CHAT CONTINUATION =====
    # Use Claude to refine feedback, then pass to NanoBanana
    if generation_mode == "chat":
        logger.info("CHAT MODE: Refining feedback with Claude, then passing to NanoBanana")
        print("\n" + "=" * 60)
        print(f"Stage 7: Grid Chat Continuation (iteration {iteration})")
        print(f"From: {refined_from} -> {iteration}")
        print(f"Feedback: {feedback[:100]}..." if len(feedback) > 100 else f"Feedback: {feedback}")
        print("=" * 60)

        # Check if previous iteration has chat history
        prev_chat_history_path = take_path / "07-grids" / refined_from / "chat_history.json"
        if not prev_chat_history_path.exists():
            return {"status": "error", "error": f"No chat history found in {refined_from}. Use 'refine' instead."}

        # Copy chat history from previous iteration
        chat_history_path = iter_path / "chat_history.json"
        import shutil
        shutil.copy(prev_chat_history_path, chat_history_path)
        logger.info(f"Copied chat history from {refined_from}")

        # Refine user feedback with Claude
        logger.info("Refining user feedback with Claude...")

        # Build grid layout description for feedback refinement
        if grid_layout == "4x4":
            grid_desc = f"Context: This is a 4x4 grid of storyboard frames ({num_scenes} scenes). The user is asking to modify specific frames."
            grid_map = """Grid layout:
Row 1: Frame 1 (top-left)     | Frame 2 (top-center-left)  | Frame 3 (top-center-right)  | Frame 4 (top-right)
Row 2: Frame 5 (upper-left)   | Frame 6 (upper-center-left)| Frame 7 (upper-center-right)| Frame 8 (upper-right)
Row 3: Frame 9 (lower-left)   | Frame 10 (lower-center-left)| Frame 11 (lower-center-right)| Frame 12 (lower-right)
Row 4: Frame 13 (bottom-left) | Frame 14 (bottom-center-left)| Frame 15 (bottom-center-right)| Frame 16 (bottom-right)"""
        elif grid_layout == "2x2":
            grid_desc = f"Context: This is a 2x2 grid of storyboard frames ({num_scenes} scenes). The user is asking to modify specific frames."
            grid_map = """Grid layout:
Row 1: Frame 1 (top-left)     | Frame 2 (top-right)
Row 2: Frame 3 (bottom-left)  | Frame 4 (bottom-right)"""
        else:
            grid_desc = f"Context: This is a 3x3 grid of storyboard frames ({num_scenes} scenes). The user is asking to modify specific frames."
            grid_map = """Grid layout:
Row 1: Frame 1 (top-left)    | Frame 2 (top-center)    | Frame 3 (top-right)
Row 2: Frame 4 (middle-left) | Frame 5 (middle-center) | Frame 6 (middle-right)
Row 3: Frame 7 (bottom-left) | Frame 8 (bottom-center) | Frame 9 (bottom-right)"""

        feedback_refine_prompt = f"""Refine this user feedback to make it a precise instruction for the image generation model.

{grid_desc}

{grid_map}

User's original feedback:
{feedback}

Your task:
1. Add position labels if frame numbers are mentioned (e.g., "Frame 7" → "Frame 7 (bottom-left)")
2. Correct any typos or unclear phrasing
3. Make the instruction more specific and verbose - tell the model exactly what to change
4. Use imperative language ("change X to Y", "move X to the right", etc.)
5. Keep it concise but extremely precise

Output only the refined feedback instruction, nothing else."""

        try:
            result = subprocess.run(
                ["claude", "--resume", session_id,
                 "--model", "opus",
                 "-p", feedback_refine_prompt],
                capture_output=True, text=True,
                cwd=str(story_path), timeout=120,
                env=claude_env()
            )

            if result.returncode == 0:
                refined_feedback = result.stdout.strip()
                logger.info(f"Feedback refined: {refined_feedback}")
                print(f"Refined feedback: {refined_feedback}")
                # Save both original and refined feedback
                (iter_path / "feedback_original.txt").write_text(feedback)
                (iter_path / "feedback_refined.txt").write_text(refined_feedback)
                prompt_for_api = refined_feedback
            else:
                error_detail = result.stderr or result.stdout
                logger.warning(f"Feedback refinement failed, using original: {error_detail}")
                prompt_for_api = feedback
        except Exception as e:
            logger.warning(f"Feedback refinement error, using original: {str(e)}")
            prompt_for_api = feedback

        # Save the feedback as the "prompt" for this iteration
        (iter_path / "nanobanana_prompt.txt").write_text(f"[Chat continuation from {refined_from}]\n\n{prompt_for_api}")

        # Create minimal generation_config for metadata
        config = {
            "generation_mode": "chat",
            "chat_from": refined_from,
            "reasoning": f"Chat continuation from {refined_from}"
        }
        (iter_path / "generation_config.json").write_text(json.dumps(config, indent=2))

        is_chat_continuation = True

    # ===== MODE 1 & 2: CLAUDE-BASED GENERATION =====
    else:
        # Build prompt
        prompt_template_path = Path(__file__).parent.parent.parent / "prompts" / "stage_grid.txt"

        if not prompt_template_path.exists():
            return {"status": "error", "error": f"Prompt template not found: {prompt_template_path}"}

        prompt = prompt_template_path.read_text()

        # Add template/references context if project has them
        project_path = take_path.parent
        refs = reference_loader.load_project_references(project_path)
        template_context = reference_loader.get_full_prompt_context(project_path)
        if template_context:
            prompt += "\n\n" + template_context
            logger.info(f"Added template context ({len(template_context)} chars)")

        # Inject creative_notes from story_pipeline.json
        story_pipeline_file = project_path / "story_pipeline.json"
        if story_pipeline_file.exists():
            try:
                sp_data = json.loads(story_pipeline_file.read_text())
                creative_notes = sp_data.get("creative_notes")
                if creative_notes:
                    prompt += (f"\n\n## CREATIVE DIRECTION\n\n"
                               f"Channel-level creative requirements:\n\n{creative_notes}\n")
                    logger.info(f"Added creative_notes ({len(creative_notes)} chars)")
            except (json.JSONDecodeError, OSError):
                pass

        # Inject interpolation frame semantics if configured
        if video_mode == "interpolation":
            prompt += """

## INTERPOLATION FRAME PAIRING

This project uses frame-pair interpolation. Frames are paired for start/end frame video generation:
- **Frame 1 + Frame 2 = Clip 1** (Frame 1 is START, Frame 2 is END)
- **Frame 3 + Frame 4 = Clip 2** (Frame 3 is START, Frame 4 is END)

**Standard mode** (default): All frames show "the moment before action begins" — anticipation states.

**Interpolation mode** (THIS project): Odd frames (1,3) = START states, even frames (2,4) = END states.

**Key requirement**: Maximize visual CONTRAST within each pair for maximum interpolation effect:
- Same composition/setting but DRAMATICALLY different state
- Same characters in same positions but transformed (calm→chaos, order→destruction, empty→crowded)
- The AI will morph between these two frames — the bigger the visual difference, the more striking the result
- Both frames in a pair MUST share the same camera angle, framing, and physical space
"""
            logger.info("Added interpolation frame pairing instructions")

        # Replace placeholders with actual V5 iteration paths
        take_name = take_path.name
        prompt = prompt.replace("{take}", take_name)
        prompt = prompt.replace("{script_path}", str(script_file.relative_to(take_path.parent)))
        prompt = prompt.replace("{assets_path}", str(assets_file.relative_to(take_path.parent)))
        prompt = prompt.replace("{char_shot_path}", str(char_shot.relative_to(take_path.parent)))

        # Add mode-specific instructions
        if generation_mode == "new":
            # Mode 1: NEW - Fresh generation
            prompt += f"\n\n## Generation Mode: NEW\n\n"
            prompt += f"This is a COMPLETELY NEW generation from scratch. You are creating something fresh.\n\n"
            prompt += f"The characteristic shot is the ONLY reference image (no previous grids).\n\n"
            if feedback:
                prompt += f"### User Feedback (as guidance)\n\n"
                prompt += f"The user provided this feedback to help avoid artifacts and inconsistencies:\n\n"
                prompt += f"{feedback}\n\n"
                prompt += f"Use this feedback as GUIDANCE only — don't over-rely on it. The scenes or requirements might have changed and may not be compatible with previous iterations. Focus on creating a cohesive new storyboard.\n"
            prompt += f"\n**generation_config.json** must have: `{{'generation_mode': 'new', 'reasoning': '...'}}`\n"

        elif generation_mode == "refine":
            # Mode 2: REFINE - Rewrite previous prompt incorporating feedback
            prompt += f"\n\n## Generation Mode: REFINE\n\n"
            prompt += f"You are REFINING the NanoBanana prompt from iteration **{refined_from}**.\n\n"
            prompt += f"**IMPORTANT**: You must KEEP everything that worked in the previous prompt (compositions, specific visual details, creative choices) and ONLY change what the feedback asks for. Do NOT rewrite from scratch.\n\n"

            # Include previous NanoBanana prompt
            prev_prompt_path = take_path / "07-grids" / refined_from / "nanobanana_prompt.txt"
            if prev_prompt_path.exists():
                prev_prompt = prev_prompt_path.read_text().strip()
                prompt += f"### Previous NanoBanana Prompt (from {refined_from})\n\n"
                prompt += f"This is the prompt that generated the previous grid. Modify it according to feedback, keeping everything else intact:\n\n"
                prompt += f"```\n{prev_prompt}\n```\n\n"
            else:
                prompt += f"(Previous prompt not found at {prev_prompt_path})\n\n"

            prompt += f"### Accumulated Feedback from All Previous Iterations\n\n"
            prompt += f"Below is ALL feedback from previous grid iterations.\n"
            prompt += f"Your task:\n"
            prompt += f"1. Read and understand ALL feedback entries\n"
            prompt += f"2. **Start from the previous NanoBanana prompt above** and modify it\n"
            prompt += f"3. If there are contradictions, **prioritize the MOST RECENT feedback** (later iterations override earlier ones)\n"
            prompt += f"4. Keep all creative details (specific objects, compositions, camera angles) that feedback does NOT mention changing\n\n"

            if all_feedback:
                for entry in all_feedback:
                    prompt += f"**Iteration {entry['iteration']}:**\n{entry['feedback']}\n\n"

            if feedback:
                prompt += f"**Current feedback (MOST RECENT - highest priority):**\n{feedback}\n\n"

            prompt += f"\n**generation_config.json** must have: `{{'generation_mode': 'refine', 'refined_from': '{refined_from}', 'feedback_count': {len(all_feedback) + (1 if feedback else 0)}, 'reasoning': '...'}}`\n"

        # Add variation mode instructions if generating multiple grids
        if variation_index is not None and total_variations is not None and total_variations > 1:
            prompt += f"\n\n## VARIATION MODE\n\n"
            prompt += f"This is variation **{variation_index + 1} of {total_variations}**.\n\n"
            prompt += f"You are generating ONE of MULTIPLE grids. Each grid must be DIFFERENT.\n\n"
            prompt += f"### Required Variations\n\n"
            prompt += f"Make THIS grid visually distinct from others by varying:\n\n"
            prompt += f"1. **Camera Position/Viewpoint**: Completely different vantage points\n"
            prompt += f"   - Scene at lake: Grid 1 from shore looking at water, Grid 2 from water looking at shore, Grid 3 aerial view\n"
            prompt += f"   - Scene in room: Grid 1 from doorway, Grid 2 from window, Grid 3 from corner\n\n"
            prompt += f"2. **Camera Distance**: Wide establishing vs medium character-focused vs close-up detail\n\n"
            prompt += f"3. **Character Positions**: Where characters are placed in the frame\n"
            prompt += f"   - Left/right of frame, foreground/background\n"
            prompt += f"   - Facing toward camera vs away vs profile\n\n"
            prompt += f"4. **Starting Moments**: Different interpretations of the moment before action\n"
            prompt += f"   - Same scene can start with different characters in focus\n"
            prompt += f"   - Different point in the 'anticipation' phase\n\n"
            prompt += f"5. **Composition Style**: Symmetric vs asymmetric, rule-of-thirds, Dutch angle\n\n"
            prompt += f"The user will review all {total_variations} grids and select the best one.\n"
            prompt += f"Think creatively about HOW to show each scene differently.\n\n"
            prompt += f"**Include in generation_config.json**: `'variation_index': {variation_index}, 'total_variations': {total_variations}`\n"

        # Add corpus reference guidance if available
        corpus_refs_file = take_path / "corpus_refs.json"
        if corpus_refs_file.exists():
            try:
                corpus_data = json.loads(corpus_refs_file.read_text())
                assignments = corpus_data.get("scene_assignments", {})
                valid_assignments = {
                    k: v for k, v in assignments.items()
                    if v and (take_path / v).exists()
                }
                if valid_assignments:
                    prompt += "\n\n## Corpus Reference Images\n\n"
                    prompt += "Reference images from the project corpus are attached, labeled by frame.\n"
                    prompt += "For each frame with a corpus reference:\n"
                    prompt += "- PRESERVE the spatial composition and central visual metaphor\n"
                    prompt += "- UPDATE characters/context to current episode subject\n"
                    prompt += "- The reference is a COMPOSITION GUIDE, not just a style filter\n\n"
                    for sn in sorted(valid_assignments.keys(), key=int):
                        ref_name = Path(valid_assignments[sn]).name
                        prompt += f"  Frame {sn}: {ref_name}\n"
                    prompt += "\n"
                    logger.info(f"Added corpus guidance for {len(valid_assignments)} frames")
            except (json.JSONDecodeError, ValueError) as e:
                logger.warning(f"Could not read corpus_refs.json for prompt: {e}")

        # Prepend output directory
        relative_iter_path = iter_path.relative_to(story_path)
        prompt = f"## Output Directory\n\nWrite ALL output files to: {relative_iter_path}/\n\n{prompt}"

        # Save prompt
        (iter_path / "prompt.txt").write_text(prompt)
        logger.info("Saved prompt")

        # ===== STEP 1: Call Claude to create NanoBanana prompt =====
        logger.info("Step 1: Calling Claude to create NanoBanana prompt...")
        print("\n" + "=" * 60)
        print(f"Stage 7: Grid Generation (iteration {iteration})")
        print(f"Session: {session_id}")
        print("=" * 60)

        try:
            result = subprocess.run(
                ["claude", "--resume", session_id,
                 "--model", "opus",
                 "--allowedTools", "Read,Write,Glob,Grep",
                 "-p", prompt],
                capture_output=True, text=True,
                cwd=str(story_path), timeout=900,  # Longer timeout for Opus
                env=claude_env()
            )

            if result.returncode != 0:
                error_detail = result.stderr or result.stdout
                error_msg = f"Claude returned code {result.returncode}: {error_detail}"
                logger.error(error_msg)
                return {"status": "error", "error": error_msg}

            logger.info("Claude completed successfully")

        except subprocess.TimeoutExpired:
            return {"status": "error", "error": "Claude timed out after 10 minutes"}
        except Exception as e:
            return {"status": "error", "error": f"Exception calling Claude: {str(e)}"}

        # Validate Claude created the prompt file
        nanobanana_prompt_file = iter_path / "nanobanana_prompt.txt"
        if not nanobanana_prompt_file.exists():
            return {"status": "error", "error": "Claude did not create nanobanana_prompt.txt"}

        nanobanana_prompt = nanobanana_prompt_file.read_text().strip()
        logger.info(f"Created nanobanana_prompt.txt ({nanobanana_prompt_file.stat().st_size} bytes)")

        # Validate generation_config.json
        generation_config_file = iter_path / "generation_config.json"
        if not generation_config_file.exists():
            return {"status": "error", "error": "Claude did not create generation_config.json"}

        logger.info(f"Created generation_config.json ({generation_config_file.stat().st_size} bytes)")

        # Extract prompt text (before metadata section)
        prompt_for_api = nanobanana_prompt
        if "---" in prompt_for_api:
            prompt_for_api = prompt_for_api.split("---")[0].strip()

        # Chat history path for this iteration (will be created fresh)
        chat_history_path = iter_path / "chat_history.json"
        is_chat_continuation = False

    # ===== STEP 2: Call NanoBanana API to generate grid image =====
    logger.info("Step 2: Calling NanoBanana API to generate grid image...")
    print("\n" + "=" * 60)
    print(f"Stage 7: Grid Generation (Step 2 - NanoBanana image)")
    print("=" * 60)

    # Load .env for API keys
    try:
        from dotenv import load_dotenv
        env_path = repo_root / ".env"
        if env_path.exists():
            load_dotenv(env_path)
    except ImportError:
        logger.warning("python-dotenv not installed, skipping .env loading")

    output_image = iter_path / "grid.png"
    reference_images = []
    ref_roles = []

    # For chat continuation, we skip reference image processing entirely
    # The model already has context from the chat history
    if generation_mode == "chat":
        # No reference images for chat mode - model has context from history
        logger.info("Chat mode: No reference images (using chat history)")
    else:
        # For NEW and REFINE modes: use characteristic shot + person reference photos
        reference_images = [str(char_shot)]
        ref_roles = [{"role": "characteristic_shot", "source": "iteration"}]

        # Add person reference photos from project references (refs loaded earlier in this block)
        character_labels = None
        if char_ref_only:
            logger.info("Char-ref-only mode: using ONLY characteristic shot as reference")
        elif refs.get("available"):
            if is_story_mode(refs):
                # Story mode: all episode characters + locations + objects
                char_data = get_episode_characters(take_path, refs)
                char_refs = select_character_refs(refs, char_data["all"])
                if char_refs:
                    # character_labels: None for char shot, then character labels
                    character_labels = [None]  # char shot is first in reference_images
                    reference_images.extend(r["path"] for r in char_refs)
                    ref_roles.extend(
                        {"role": "character_ref", "asset_id": r["character_id"], "source": "reference"}
                        for r in char_refs
                    )
                    character_labels.extend(
                        {"character_id": r["character_id"], "label": r.get("description", "")}
                        for r in char_refs
                    )
                    logger.info(f"Story mode: added {len(char_refs)} character reference sheets")

                # Add location refs with frame index fallback
                from lib.frame_index import load_frame_index
                frame_index = load_frame_index(project_path)

                loc_data = get_episode_locations(take_path, refs)
                loc_refs = select_asset_refs_with_fallback(
                    refs, loc_data["all"], "locations", take_path, frame_index
                )
                remaining = NANOBANANA_MAX_IMAGES - len(reference_images)
                if loc_refs and remaining > 0:
                    loc_refs = loc_refs[:remaining]
                    reference_images.extend(r["path"] for r in loc_refs)
                    ref_roles.extend(
                        {"role": "location_ref", "asset_id": r["asset_id"],
                         "source": "frame" if r["source"] == "frame" else "reference"}
                        for r in loc_refs
                    )
                    if character_labels is not None:
                        character_labels.extend(
                            {"type": "frame_reference", "asset_id": r["asset_id"],
                             "label": r.get("description", "")}
                            if r["source"] == "frame" else None
                            for r in loc_refs
                        )
                    logger.info(f"Story mode: added {len(loc_refs)} location references "
                               f"({sum(1 for r in loc_refs if r['source'] == 'frame')} from frame index)")

                # Add object refs with frame index fallback
                obj_data = get_episode_objects(take_path, refs)
                obj_refs = select_asset_refs_with_fallback(
                    refs, obj_data["all"], "objects", take_path, frame_index
                )
                remaining = NANOBANANA_MAX_IMAGES - len(reference_images)
                if obj_refs and remaining > 0:
                    obj_refs = obj_refs[:remaining]
                    reference_images.extend(r["path"] for r in obj_refs)
                    ref_roles.extend(
                        {"role": "object_ref", "asset_id": r["asset_id"],
                         "source": "frame" if r["source"] == "frame" else "reference"}
                        for r in obj_refs
                    )
                    if character_labels is not None:
                        character_labels.extend(
                            {"type": "frame_reference", "asset_id": r["asset_id"],
                             "label": r.get("description", "")}
                            if r["source"] == "frame" else None
                            for r in obj_refs
                        )
                    logger.info(f"Story mode: added {len(obj_refs)} object references "
                               f"({sum(1 for r in obj_refs if r['source'] == 'frame')} from frame index)")
            else:
                # Product mode: single "main" person
                main_person = refs.get("people", {}).get("main", {})
                person_photos = main_person.get("photos", [])
                if person_photos:
                    photos_to_add = person_photos[:3]
                    reference_images.extend(photos_to_add)
                    ref_roles.extend(
                        {"role": "person_ref", "source": "reference"}
                        for _ in photos_to_add
                    )
                    logger.info(f"Added {len(photos_to_add)} person reference photos")

        logger.info(f"Reference images: {len(reference_images)} total (char shot + person refs)")

        # Load corpus references if available (generic feature)
        corpus_refs_map = {}  # scene_num -> abs_image_path
        corpus_refs_file = take_path / "corpus_refs.json"
        if corpus_refs_file.exists():
            try:
                corpus_data = json.loads(corpus_refs_file.read_text())
                for scene_str, ref_path in corpus_data.get("scene_assignments", {}).items():
                    if ref_path:
                        abs_path = take_path / ref_path
                        if abs_path.exists():
                            corpus_refs_map[int(scene_str)] = str(abs_path)
                if corpus_refs_map:
                    logger.info(f"Loaded {len(corpus_refs_map)} corpus refs from corpus_refs.json")
                    for scene_num in sorted(corpus_refs_map.keys()):
                        reference_images.append(corpus_refs_map[scene_num])
                        ref_roles.append({"role": "corpus_ref", "scene": scene_num, "source": "corpus"})
                        if character_labels is not None:
                            character_labels.append(
                                {"type": "corpus_ref", "scene": scene_num,
                                 "label": ""}
                            )
            except (json.JSONDecodeError, ValueError) as e:
                logger.warning(f"Could not load corpus_refs.json: {e}")

    # Prepend [REFERENCE IMAGES...] prefix for reference image guidance + sharpness
    if reference_images and generation_mode != "chat":
        refs_path = project_path / "references" if refs.get("available") else None
        index = refs.get("index")
        analyze_prefix = build_analyze_prefix(
            index, reference_images, refs_path,
            character_labels=character_labels
        )
        if analyze_prefix:
            prompt_for_api = analyze_prefix + prompt_for_api
            # Save the full prompt (with reference prefix) for auditability
            (iter_path / "nanobanana_prompt.txt").write_text(prompt_for_api)
            logger.info(f"Added [REFERENCE IMAGES...] prefix ({len(analyze_prefix)} chars)")

    # Scale down references if requested
    if scale_refs and reference_images:
        from lib import scale_down_references
        scaled_dir = iter_path / "scaled_refs"
        reference_images = scale_down_references(
            reference_images, scaled_dir, max_pixels=2048
        )
        logger.info(f"Scaled references to 2K in {scaled_dir}")

    try:
        from nanobanana_text2img import generate_image_chat

        # Use chat mode for generation
        result_path = generate_image_chat(
            prompt=prompt_for_api,
            output_path=str(output_image),
            history_path=str(chat_history_path),
            aspect_ratio="9:16",
            image_size="4K",
            reference_image_paths=reference_images or None,
            is_refinement=is_chat_continuation
        )

        logger.info(f"Image generated: {result_path}")
        output_image = Path(result_path)

    except Exception as e:
        error_msg = f"NanoBanana image generation failed: {str(e)}"
        logger.error(error_msg)
        return {"status": "error", "error": error_msg}

    # User message
    user_message_file = iter_path / "user_message.txt"
    user_message = (user_message_file.read_text().strip()
                    if user_message_file.exists()
                    else "Grid image generated. Please review.")

    # Save metadata (V5 format)
    metadata = {
        "stage": "07-grids",
        "iteration": iteration,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "based_on": inputs,
        "refined_from": refined_from,
        "rating": None,
        "result": f"07-grids/{iteration}/{output_image.name}",
        "generation_mode": generation_mode,
        "model_prompt": "claude-opus",
        "model_image": "nanobanana-gemini3pro",
        "reference_images_sent": build_reference_record(
            reference_images or [], take_path.parent, ref_roles
        ),
        "session_id": session_id,
        "image_file": output_image.name,
        "aspect_ratio": "9:16",
        "grid_layout": grid_layout,
        "scene_count": num_scenes,
        "characteristic_ref": str(char_shot),
        "chat_history": "chat_history.json" if chat_history_path.exists() else None,
    }

    # Add variation info if batch generation
    if variation_index is not None:
        metadata["variation_index"] = variation_index
    if total_variations is not None:
        metadata["total_variations"] = total_variations

    save_metadata(iter_path / "metadata.json", metadata)
    logger.info("Saved metadata")

    # Log event
    log_event(
        story_path, "stage_complete",
        {"action": "grid_generated", "iteration": iteration,
         "result": f"07-grids/{iteration}/{output_image.name}"},
        take=take_path.name, stage="07-grids", iteration=iteration
    )

    print(f"\nGenerated: 07-grids/{iteration}/{output_image.name}")

    return {
        "status": "success",
        "iteration": iteration,
        "result": f"07-grids/{iteration}/{output_image.name}",
        "user_message": user_message,
    }
