#!/usr/bin/env python3
"""
Stage S5: Asset Reference Sheets

Per-asset iteration model. Each character, location, and object gets
independent iterations that can be individually refined, regenerated,
or chat-refined via NanoBanana.

Directory structure:
    S5-assets/
        assets_plan.json          # Asset plan + depends_on (editable)
        selections.json           # Per-asset selected iterations
        characters/{id}/
            001/                  # Iteration (self-contained)
                prompt.txt        # NanoBanana prompt
                description.md    # Character description (characters only)
                sheet.jpg         # Generated reference sheet
                chat_history.json # NanoBanana conversation
                metadata.json
        locations/{id}/
            001/
                prompt.txt, sheet.jpg, chat_history.json, metadata.json
        objects/{id}/
            001/
                prompt.txt, sheet.jpg, chat_history.json, metadata.json

Dependencies:
    assets_plan.json entries can have "depends_on": ["other_asset_id"]
    When generating, dependency sheets are passed as additional NanoBanana
    reference images alongside style_board and title_artwork.
"""

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

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 claude_env, save_metadata, ensure_dir
from lib.path_utils_v3 import get_next_iteration
from lib.reference_tracking import build_reference_record


ASSET_TYPES = ["characters", "locations", "objects"]
# Generate objects first (standalone), then characters, then locations
# (which may depend on objects/characters)
GENERATION_ORDER = ["objects", "characters", "locations"]
ASSET_IMAGE_NAME = "sheet.jpg"


def execute(story_path: Path, inputs: dict,
            feedback: str = None, refined_from: str = None) -> dict:
    """
    Run Stage S5: Create assets plan + generate all reference sheets.

    First call: Claude creates assets_plan.json with prompts for each asset.
    Then NanoBanana generates reference sheets for all assets.
    Subsequent calls resume from where generation left off.
    """
    stage_dir = story_path / "S5-assets"
    ensure_dir(stage_dir)

    print(f"\nStage S5: Asset Reference Sheets")

    # Resolve inputs
    bible_iter = inputs.get("S2-bible")
    style_iter = inputs.get("S3-style")
    if not bible_iter:
        return {"status": "error", "error": "S2-bible not in inputs"}
    if not style_iter:
        return {"status": "error", "error": "S3-style not in inputs"}

    bible_path = f"S2-bible/{bible_iter}/story_bible.json"
    style_path = f"S3-style/{style_iter}/style.json"

    if not (story_path / bible_path).exists():
        return {"status": "error", "error": f"Story bible not found: {bible_path}"}
    if not (story_path / style_path).exists():
        return {"status": "error", "error": f"Style not found: {style_path}"}

    # Load session
    session_file = story_path / "session.txt"
    if not session_file.exists():
        return {"status": "error", "error": "session.txt not found"}
    session_id = session_file.read_text().strip()

    # Check if plan already exists (resume from failed generation)
    plan_file = stage_dir / "assets_plan.json"
    if not plan_file.exists():
        # Step 1: Call Claude to create assets plan + prompts
        prompt_path = (Path(__file__).parent.parent.parent
                       / "prompts" / "story" / "stage_s5_assets.txt")
        if not prompt_path.exists():
            return {"status": "error",
                    "error": f"Prompt template not found: {prompt_path}"}

        prompt = prompt_path.read_text()
        prompt = prompt.replace("{bible_path}", bible_path)
        prompt = prompt.replace("{style_path}", style_path)

        if feedback:
            prompt += f"\n\n## User Feedback\n\n{feedback}\n"

        # Output directory
        relative_path = stage_dir.relative_to(story_path)
        prompt = (f"## Output Directory\n\nWrite ALL output files to: "
                  f"{relative_path}/\n\n{prompt}")
        (stage_dir / "prompt.txt").write_text(prompt)

        print(f"Session: {session_id}")
        print("Calling Claude to plan assets and create NanoBanana prompts...")

        try:
            result = subprocess.run(
                ["claude", "--resume", session_id,
                 "--allowedTools", "Read,Write,Glob,Grep",
                 "-p", prompt],
                capture_output=True, text=True,
                cwd=str(story_path), timeout=600,
                env=claude_env()
            )
            if result.returncode != 0:
                error_detail = result.stderr or result.stdout
                return {"status": "error",
                        "error": f"Claude returned code {result.returncode}: "
                                 f"{error_detail}"}
        except subprocess.TimeoutExpired:
            return {"status": "error", "error": "Claude timed out"}

        if not plan_file.exists():
            return {"status": "error",
                    "error": "assets_plan.json not created by Claude"}
    else:
        print("Assets plan exists, generating remaining sheets...")

    # Step 2: Read plan and generate sheets
    plan = json.loads(plan_file.read_text())
    selections = load_selections(stage_dir)

    # Style reference images
    style_refs = _get_style_references(story_path)

    generated = []
    failed = []

    # Generate in dependency order: objects first, then characters, then locations
    for asset_type in GENERATION_ORDER:
        for asset_info in plan.get(asset_type, []):
            asset_id = asset_info["id"]
            asset_dir = stage_dir / asset_type / asset_id

            # Skip if already has an iteration
            existing_iters = _list_asset_iterations(asset_dir)
            if existing_iters:
                sel = selections.get(asset_type, {}).get(asset_id)
                status = f"selected={sel}" if sel else "generated"
                print(f"  {asset_type}/{asset_id}: {status} (skip)")
                generated.append(f"{asset_type}/{asset_id}")
                continue

            # Read prompt from staging location (Claude wrote it here)
            prompt_rel = asset_info.get(
                "prompt_file", f"{asset_type}/{asset_id}/prompt.txt")
            prompt_file = stage_dir / prompt_rel
            if not prompt_file.exists():
                print(f"  Warning: No prompt file for {asset_type}/{asset_id}")
                failed.append(f"{asset_type}/{asset_id}")
                continue

            # Read description (characters only)
            desc_rel = asset_info.get("description_file")
            desc_file = stage_dir / desc_rel if desc_rel else None

            # Resolve dependency references
            dep_refs = _resolve_dependency_refs(stage_dir, plan, asset_info)

            # Combine style + dependency refs
            all_refs = (style_refs or []) + dep_refs
            ref_images = all_refs if all_refs else None

            print(f"Generating sheet for {asset_type}/{asset_id}...")
            if dep_refs:
                dep_ids = asset_info.get("depends_on", [])
                print(f"  Dependencies: {', '.join(dep_ids)} "
                      f"({len(dep_refs)} ref images)")

            result_path = _generate_asset_sheet(
                asset_dir=asset_dir,
                prompt_file=prompt_file,
                desc_file=desc_file,
                ref_images=ref_images,
                asset_type=asset_type,
                asset_id=asset_id,
                story_path=story_path,
            )

            if result_path:
                iteration = result_path.parent.name
                generated.append(f"{asset_type}/{asset_id}")

                # Auto-select first iteration
                if asset_type not in selections:
                    selections[asset_type] = {}
                selections[asset_type][asset_id] = iteration

                # Copy to references
                _copy_to_references(
                    story_path, asset_type, asset_id,
                    result_path)

                print(f"  -> {asset_type}/{asset_id}/{iteration}/sheet.jpg")
            else:
                failed.append(f"{asset_type}/{asset_id}")
                print(f"  Warning: Generation failed for "
                      f"{asset_type}/{asset_id}")

    # Save selections
    save_selections(stage_dir, selections)

    # User message
    user_msg_file = stage_dir / "user_message.txt"
    user_message = (user_msg_file.read_text().strip()
                    if user_msg_file.exists()
                    else f"Generated {len(generated)} asset sheets.")
    if failed:
        user_message += f"\nFailed: {', '.join(failed)}"

    # Save metadata
    metadata = {
        "stage": "S5-assets",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "based_on": inputs,
        "session_id": session_id,
        "generated": generated,
        "failed": failed,
    }
    save_metadata(stage_dir / "metadata.json", metadata)

    print(f"\nGenerated: {len(generated)} asset sheets")
    if failed:
        print(f"Failed: {len(failed)} ({', '.join(failed)})")

    return {
        "status": "success",
        "generated": generated,
        "failed": failed,
        "user_message": user_message,
    }


def execute_single(story_path: Path, asset_id: str,
                   mode: str, feedback: str = None) -> dict:
    """
    Operate on a single asset: refine, regenerate, or chat.

    Args:
        asset_id: Asset ID (e.g., "tom_morrison", "scorpion_desert")
        mode: "refine", "regenerate", or "chat"
        feedback: Feedback/prompt text
    """
    stage_dir = story_path / "S5-assets"
    plan_file = stage_dir / "assets_plan.json"

    if not plan_file.exists():
        return {"status": "error",
                "error": "No assets plan. Run S5-assets first."}

    plan = json.loads(plan_file.read_text())

    # Find asset in plan
    asset_type, asset_info = _find_asset(plan, asset_id)
    if not asset_type:
        available = _list_all_asset_ids(plan)
        return {"status": "error",
                "error": f"Asset '{asset_id}' not found. "
                         f"Available: {', '.join(available)}"}

    asset_dir = stage_dir / asset_type / asset_id
    selections = load_selections(stage_dir)

    # Resolve all references: style + dependencies
    style_refs = _get_style_references(story_path)
    dep_refs = _resolve_dependency_refs(stage_dir, plan, asset_info)
    all_refs = (style_refs or []) + dep_refs
    ref_images = all_refs if all_refs else None

    # Find current selected iteration
    current_sel = selections.get(asset_type, {}).get(asset_id)
    if not current_sel:
        existing = _list_asset_iterations(asset_dir)
        current_sel = existing[-1] if existing else None

    if mode == "refine":
        if not feedback:
            return {"status": "error",
                    "error": "Refine requires feedback text"}

        if not current_sel:
            return {"status": "error",
                    "error": f"No iterations for {asset_id}. "
                             "Run S5-assets first."}

        # Read prompt from current iteration and append feedback
        src_dir = asset_dir / current_sel
        prompt_file = src_dir / "prompt.txt"
        if not prompt_file.exists():
            return {"status": "error",
                    "error": f"No prompt.txt in {asset_id}/{current_sel}"}

        base_prompt = prompt_file.read_text().strip()
        modified_prompt = f"{base_prompt}\n\nREFINEMENT: {feedback}"

        # Create new iteration
        iteration = get_next_iteration(asset_dir)
        iter_dir = asset_dir / iteration
        ensure_dir(iter_dir)

        # Write modified prompt
        (iter_dir / "prompt.txt").write_text(modified_prompt)

        # Copy description.md from source iteration (characters)
        src_desc = src_dir / "description.md"
        if src_desc.exists():
            shutil.copy2(src_desc, iter_dir / "description.md")

        print(f"Refining {asset_type}/{asset_id} "
              f"(iteration {iteration})...")
        print(f"  Feedback: {feedback}")
        if dep_refs:
            print(f"  Dependencies: {len(dep_refs)} ref images")

        result_path = _generate_from_prompt(
            prompt_text=modified_prompt,
            output_path=iter_dir / ASSET_IMAGE_NAME,
            chat_history_path=iter_dir / "chat_history.json",
            ref_images=ref_images,
        )

        if not result_path:
            shutil.rmtree(iter_dir, ignore_errors=True)
            return {"status": "error",
                    "error": "NanoBanana generation failed"}

        ref_roles = _build_asset_ref_roles(ref_images) if ref_images else []
        save_metadata(iter_dir / "metadata.json", {
            "asset_type": asset_type,
            "asset_id": asset_id,
            "iteration": iteration,
            "mode": "refine",
            "feedback": feedback,
            "refined_from": current_sel,
            "created_at": datetime.now(timezone.utc).isoformat(),
            "reference_images_sent": build_reference_record(
                ref_images or [], story_path, ref_roles
            ),
        })

        # Auto-select new iteration
        if asset_type not in selections:
            selections[asset_type] = {}
        selections[asset_type][asset_id] = iteration
        save_selections(stage_dir, selections)

        _copy_to_references(
            story_path, asset_type, asset_id, result_path)

        return {
            "status": "success",
            "asset": f"{asset_type}/{asset_id}",
            "iteration": iteration,
            "mode": "refine",
        }

    elif mode == "regenerate":
        if not current_sel:
            return {"status": "error",
                    "error": f"No iterations for {asset_id}. "
                             "Run S5-assets first."}

        # Read prompt from current iteration
        src_dir = asset_dir / current_sel
        prompt_file = src_dir / "prompt.txt"
        if not prompt_file.exists():
            return {"status": "error",
                    "error": f"No prompt.txt in {asset_id}/{current_sel}"}

        prompt_text = prompt_file.read_text().strip()
        if feedback:
            prompt_text = f"{prompt_text}\n\nADDITIONAL DIRECTION: {feedback}"

        # Create new iteration
        iteration = get_next_iteration(asset_dir)
        iter_dir = asset_dir / iteration
        ensure_dir(iter_dir)

        (iter_dir / "prompt.txt").write_text(prompt_text)

        # Copy description.md from source iteration (characters)
        src_desc = src_dir / "description.md"
        if src_desc.exists():
            shutil.copy2(src_desc, iter_dir / "description.md")

        print(f"Regenerating {asset_type}/{asset_id} "
              f"(iteration {iteration})...")
        if dep_refs:
            print(f"  Dependencies: {len(dep_refs)} ref images")

        result_path = _generate_from_prompt(
            prompt_text=prompt_text,
            output_path=iter_dir / ASSET_IMAGE_NAME,
            chat_history_path=iter_dir / "chat_history.json",
            ref_images=ref_images,
        )

        if not result_path:
            shutil.rmtree(iter_dir, ignore_errors=True)
            return {"status": "error",
                    "error": "NanoBanana generation failed"}

        ref_roles = _build_asset_ref_roles(ref_images) if ref_images else []
        save_metadata(iter_dir / "metadata.json", {
            "asset_type": asset_type,
            "asset_id": asset_id,
            "iteration": iteration,
            "mode": "regenerate",
            "feedback": feedback,
            "refined_from": current_sel,
            "created_at": datetime.now(timezone.utc).isoformat(),
            "reference_images_sent": build_reference_record(
                ref_images or [], story_path, ref_roles
            ),
        })

        # Auto-select new iteration
        if asset_type not in selections:
            selections[asset_type] = {}
        selections[asset_type][asset_id] = iteration
        save_selections(stage_dir, selections)

        _copy_to_references(
            story_path, asset_type, asset_id, result_path)

        return {
            "status": "success",
            "asset": f"{asset_type}/{asset_id}",
            "iteration": iteration,
            "mode": "regenerate",
        }

    elif mode == "chat":
        if not feedback:
            return {"status": "error",
                    "error": "Chat requires a prompt"}

        if not current_sel:
            return {"status": "error",
                    "error": f"No iterations for {asset_id}. "
                             "Run S5-assets first."}

        src_dir = asset_dir / current_sel
        chat_history = src_dir / "chat_history.json"

        if not chat_history.exists():
            return {"status": "error",
                    "error": f"No chat history in "
                             f"{asset_id}/{current_sel}"}

        # Create new iteration
        iteration = get_next_iteration(asset_dir)
        iter_dir = asset_dir / iteration
        ensure_dir(iter_dir)

        # Copy files from source iteration
        for item in src_dir.iterdir():
            if item.is_file():
                shutil.copy2(item, iter_dir / item.name)

        new_chat_history = iter_dir / "chat_history.json"
        new_output = iter_dir / ASSET_IMAGE_NAME

        print(f"Chat-refining {asset_type}/{asset_id} "
              f"(iteration {iteration})...")
        print(f"  Prompt: {feedback}")
        if dep_refs:
            print(f"  Dependencies: {len(dep_refs)} ref images")

        try:
            from nanobanana_text2img import generate_image_chat

            result_path = generate_image_chat(
                prompt=feedback,
                output_path=str(new_output),
                history_path=str(new_chat_history),
                aspect_ratio="9:16",
                image_size="4K",
                reference_image_paths=ref_images,
                is_refinement=True,
            )
            result_path = Path(result_path)
        except Exception as e:
            shutil.rmtree(iter_dir, ignore_errors=True)
            return {"status": "error",
                    "error": f"NanoBanana chat failed: {e}"}

        ref_roles = _build_asset_ref_roles(ref_images) if ref_images else []
        save_metadata(iter_dir / "metadata.json", {
            "asset_type": asset_type,
            "asset_id": asset_id,
            "iteration": iteration,
            "mode": "chat",
            "chat_prompt": feedback,
            "refined_from": current_sel,
            "created_at": datetime.now(timezone.utc).isoformat(),
            "reference_images_sent": build_reference_record(
                ref_images or [], story_path, ref_roles
            ),
        })

        # Auto-select new iteration
        if asset_type not in selections:
            selections[asset_type] = {}
        selections[asset_type][asset_id] = iteration
        save_selections(stage_dir, selections)

        _copy_to_references(
            story_path, asset_type, asset_id, result_path)

        return {
            "status": "success",
            "asset": f"{asset_type}/{asset_id}",
            "iteration": iteration,
            "mode": "chat",
        }

    else:
        return {"status": "error", "error": f"Unknown mode: {mode}"}


def accept_all(story_path: Path) -> tuple:
    """Accept all assets — select latest iteration for any unselected.

    Returns:
        (success: bool, message: str)
    """
    stage_dir = story_path / "S5-assets"
    plan_file = stage_dir / "assets_plan.json"

    if not plan_file.exists():
        return False, "No assets plan. Run S5-assets first."

    plan = json.loads(plan_file.read_text())
    selections = load_selections(stage_dir)

    accepted = []
    missing = []

    for asset_type in ASSET_TYPES:
        for asset_info in plan.get(asset_type, []):
            asset_id = asset_info["id"]
            asset_dir = stage_dir / asset_type / asset_id

            current_sel = selections.get(asset_type, {}).get(asset_id)
            if current_sel:
                accepted.append(f"{asset_type}/{asset_id}: {current_sel}")
                continue

            # Auto-select latest
            iters = _list_asset_iterations(asset_dir)
            if iters:
                latest = iters[-1]
                if asset_type not in selections:
                    selections[asset_type] = {}
                selections[asset_type][asset_id] = latest
                accepted.append(
                    f"{asset_type}/{asset_id}: {latest} (auto-selected)")

                # Copy to references
                result_path = asset_dir / latest / ASSET_IMAGE_NAME
                if result_path.exists():
                    _copy_to_references(
                        story_path, asset_type, asset_id, result_path)
            else:
                missing.append(f"{asset_type}/{asset_id}")

    save_selections(stage_dir, selections)

    if missing:
        msg = f"Cannot accept: missing sheets for {', '.join(missing)}"
        return False, msg

    msg = (f"Accepted {len(accepted)} assets:\n"
           + "\n".join(f"  {a}" for a in accepted))
    return True, msg


def get_asset_status(story_path: Path) -> str:
    """Get formatted status of all assets."""
    stage_dir = story_path / "S5-assets"
    plan_file = stage_dir / "assets_plan.json"

    if not plan_file.exists():
        return "S5-assets: No plan yet"

    plan = json.loads(plan_file.read_text())
    selections = load_selections(stage_dir)

    lines = ["S5 Asset Reference Sheets:", ""]

    total = 0
    selected = 0

    for asset_type in ASSET_TYPES:
        assets = plan.get(asset_type, [])
        if not assets:
            continue

        lines.append(f"  {asset_type.title()}:")
        for asset_info in assets:
            asset_id = asset_info["id"]
            asset_name = asset_info.get("name", asset_id)
            asset_dir = stage_dir / asset_type / asset_id
            iters = _list_asset_iterations(asset_dir)
            sel = selections.get(asset_type, {}).get(asset_id)
            deps = asset_info.get("depends_on", [])

            total += 1
            if sel:
                selected += 1

            n = len(iters)
            iter_str = f"{n} iteration{'s' if n != 1 else ''}"
            sel_str = f"selected={sel}" if sel else "no selection"
            dep_str = f" deps=[{','.join(deps)}]" if deps else ""
            lines.append(
                f"    {asset_id} ({asset_name}): "
                f"{iter_str}, {sel_str}{dep_str}")
        lines.append("")

    lines.append(f"  Total: {selected}/{total} selected")

    return "\n".join(lines)


def select_asset(story_path: Path, asset_id: str,
                 iteration: str) -> tuple:
    """Select a specific iteration for an asset.

    Returns:
        (success: bool, message: str)
    """
    stage_dir = story_path / "S5-assets"
    plan_file = stage_dir / "assets_plan.json"

    if not plan_file.exists():
        return False, "No assets plan."

    plan = json.loads(plan_file.read_text())
    asset_type, asset_info = _find_asset(plan, asset_id)

    if not asset_type:
        available = _list_all_asset_ids(plan)
        return False, (f"Asset '{asset_id}' not found. "
                       f"Available: {', '.join(available)}")

    asset_dir = stage_dir / asset_type / asset_id
    iter_dir = asset_dir / iteration

    if not iter_dir.exists():
        iters = _list_asset_iterations(asset_dir)
        avail = ', '.join(iters) if iters else 'none'
        return False, (f"Iteration {iteration} not found. "
                       f"Available: {avail}")

    selections = load_selections(stage_dir)
    if asset_type not in selections:
        selections[asset_type] = {}
    selections[asset_type][asset_id] = iteration
    save_selections(stage_dir, selections)

    # Copy to references
    result_path = iter_dir / ASSET_IMAGE_NAME
    if result_path.exists():
        _copy_to_references(
            story_path, asset_type, asset_id, result_path)

    return True, f"Selected: {asset_type}/{asset_id}/{iteration}"


# ============================================================
# Helpers
# ============================================================


def load_selections(stage_dir: Path) -> dict:
    """Load per-asset selections."""
    sel_file = stage_dir / "selections.json"
    if sel_file.exists():
        return json.loads(sel_file.read_text())
    return {}


def save_selections(stage_dir: Path, selections: dict) -> None:
    """Save per-asset selections."""
    sel_file = stage_dir / "selections.json"
    with open(sel_file, "w") as f:
        json.dump(selections, f, indent=2)
        f.write("\n")


def _find_asset(plan: dict, asset_id: str) -> tuple:
    """Find an asset by ID across all types.

    Returns:
        (asset_type, asset_info) or (None, None)
    """
    for asset_type in ASSET_TYPES:
        for info in plan.get(asset_type, []):
            if info["id"] == asset_id:
                return asset_type, info
    return None, None


def _list_all_asset_ids(plan: dict) -> list:
    """List all asset IDs from plan."""
    ids = []
    for asset_type in ASSET_TYPES:
        for info in plan.get(asset_type, []):
            ids.append(info["id"])
    return ids


def _list_asset_iterations(asset_dir: Path) -> list:
    """List iteration directories for an asset (sorted)."""
    if not asset_dir.exists():
        return []
    iters = []
    for d in sorted(asset_dir.iterdir()):
        if d.is_dir() and d.name.isdigit():
            iters.append(d.name)
    return iters


def _get_style_references(story_path: Path) -> list | None:
    """Get style board + title artwork paths for NanoBanana reference."""
    refs = []
    style_board = story_path / "references" / "style" / "style_board.jpg"
    if style_board.exists():
        refs.append(str(style_board))
    title_artwork = story_path / "references" / "style" / "title_artwork.jpg"
    if title_artwork.exists():
        refs.append(str(title_artwork))
    return refs if refs else None


def _resolve_dependency_refs(stage_dir: Path, plan: dict,
                             asset_info: dict) -> list:
    """Resolve depends_on to sheet.jpg paths of dependency assets."""
    depends_on = asset_info.get("depends_on", [])
    if not depends_on:
        return []

    selections = load_selections(stage_dir)
    refs = []

    for dep_id in depends_on:
        dep_type, dep_info = _find_asset(plan, dep_id)
        if not dep_type:
            print(f"  Warning: dependency '{dep_id}' not found in plan")
            continue

        # Use selected iteration, or latest available
        sel = selections.get(dep_type, {}).get(dep_id)
        if not sel:
            dep_dir = stage_dir / dep_type / dep_id
            iters = _list_asset_iterations(dep_dir)
            sel = iters[-1] if iters else None

        if sel:
            sheet = stage_dir / dep_type / dep_id / sel / ASSET_IMAGE_NAME
            if sheet.exists():
                refs.append(str(sheet))
                print(f"  Dep ref: {dep_type}/{dep_id}/{sel}/sheet.jpg")

    return refs


def _generate_asset_sheet(asset_dir: Path, prompt_file: Path,
                          desc_file: Path | None,
                          ref_images: list | None,
                          asset_type: str, asset_id: str,
                          story_path: Path = None) -> Path | None:
    """Generate initial sheet for an asset.

    Moves prompt.txt (and description.md for characters) from the staging
    location (asset dir) into the iteration directory. Returns sheet path
    or None on failure.
    """
    iteration = get_next_iteration(asset_dir)
    iter_dir = asset_dir / iteration
    ensure_dir(iter_dir)

    # Move prompt from staging to iteration
    prompt_text = prompt_file.read_text().strip()
    (iter_dir / "prompt.txt").write_text(prompt_text)
    # Remove staging file (now lives in iteration)
    if prompt_file.parent == asset_dir:
        prompt_file.unlink(missing_ok=True)

    # Move description.md for characters
    if desc_file and desc_file.exists():
        shutil.move(str(desc_file), str(iter_dir / "description.md"))

    # Build ref_roles for tracking
    ref_roles = _build_asset_ref_roles(ref_images) if ref_images else []

    result_path = _generate_from_prompt(
        prompt_text=prompt_text,
        output_path=iter_dir / ASSET_IMAGE_NAME,
        chat_history_path=iter_dir / "chat_history.json",
        ref_images=ref_images,
    )

    if result_path:
        meta = {
            "asset_type": asset_type,
            "asset_id": asset_id,
            "iteration": iteration,
            "mode": "initial",
            "created_at": datetime.now(timezone.utc).isoformat(),
            "reference_images_sent": build_reference_record(
                ref_images or [], story_path or asset_dir.parent.parent.parent, ref_roles
            ),
        }
        save_metadata(iter_dir / "metadata.json", meta)
        return result_path

    # Clean up on failure
    shutil.rmtree(iter_dir, ignore_errors=True)
    return None


def _build_asset_ref_roles(ref_images: list) -> list:
    """Build ref_roles list for asset reference images.

    Asset refs follow a predictable pattern: style_board, title_artwork,
    then dependency_ref images.
    """
    roles = []
    for img_path in ref_images:
        p = Path(img_path)
        name = p.stem
        parent = p.parent.name

        if "style_board" in name:
            roles.append({"role": "style_board", "source": "reference"})
        elif "title_artwork" in name:
            roles.append({"role": "title_artwork", "source": "reference"})
        elif parent == "style":
            roles.append({"role": "style_board", "source": "reference"})
        else:
            # Dependency ref (sheet from another asset)
            roles.append({"role": "dependency_ref", "asset_id": p.parent.parent.name,
                           "source": "iteration"})
    return roles


def _generate_from_prompt(prompt_text: str, output_path: Path,
                          chat_history_path: Path,
                          ref_images: list | None = None) -> Path | None:
    """Generate an image via NanoBanana API. Returns Path or None."""
    try:
        from dotenv import load_dotenv
        env_path = Path(__file__).parent.parent.parent.parent / ".env"
        if env_path.exists():
            load_dotenv(env_path)
    except ImportError:
        pass

    if not prompt_text:
        return None

    try:
        from nanobanana_text2img import generate_image_chat

        result_path = generate_image_chat(
            prompt=prompt_text,
            output_path=str(output_path),
            history_path=str(chat_history_path),
            aspect_ratio="9:16",
            image_size="4K",
            reference_image_paths=ref_images,
        )
        print(f"  Generated: {Path(result_path).name}")
        return Path(result_path)

    except Exception as e:
        print(f"  NanoBanana error: {e}")
        return None


def list_missing_assets(story_path: Path) -> str:
    """Scan episodes' 04-assets files, compare against references/ directory.

    Returns formatted text showing assets that appear in episode scenes but
    have no reference sheet in the references/ directory.
    """
    stage_dir = story_path / "S5-assets"
    plan_file = stage_dir / "assets_plan.json"
    plan = json.loads(plan_file.read_text()) if plan_file.exists() else {}
    plan_ids = set(_list_all_asset_ids(plan))

    # Collect all assets from episode 04-assets files
    # Structure: {asset_type: {asset_id: {"description": str, "episodes": {ep: [scenes]}}}}
    discovered = {"characters": {}, "locations": {}, "objects": {}}

    # Find episode directories (3-digit pattern)
    for ep_dir in sorted(story_path.iterdir()):
        if not ep_dir.is_dir() or not ep_dir.name.isdigit() or len(ep_dir.name) != 3:
            continue
        ep_num = ep_dir.name

        # Find latest iteration of 04-assets
        assets_stage = ep_dir / "04-assets"
        if not assets_stage.exists():
            continue

        # Get latest iteration
        iters = sorted(
            [d for d in assets_stage.iterdir() if d.is_dir() and d.name.isdigit()],
            key=lambda d: d.name
        )
        if not iters:
            continue

        assets_file = iters[-1] / "assets.json"
        if not assets_file.exists():
            continue

        try:
            data = json.loads(assets_file.read_text())
        except (json.JSONDecodeError, OSError):
            continue

        for asset_type in ["characters", "locations", "objects"]:
            for asset in data.get(asset_type, []):
                aid = asset.get("id")
                if not aid:
                    continue

                if aid not in discovered[asset_type]:
                    desc = asset.get("description", asset.get("appearance", ""))
                    discovered[asset_type][aid] = {
                        "description": desc,
                        "name": asset.get("name", aid),
                        "episodes": {},
                    }

                scenes = asset.get("appears_in", [])
                discovered[asset_type][aid]["episodes"][ep_num] = scenes

    # Determine which reference files exist
    refs_dir = story_path / "references"
    existing_refs = set()

    # Characters: references/people/{id}/
    people_dir = refs_dir / "people"
    if people_dir.exists():
        for d in people_dir.iterdir():
            if d.is_dir():
                existing_refs.add(("characters", d.name))

    # Locations: references/locations/{id}.jpg
    loc_dir = refs_dir / "locations"
    if loc_dir.exists():
        for f in loc_dir.iterdir():
            if f.is_file() and f.suffix in (".jpg", ".png"):
                existing_refs.add(("locations", f.stem))

    # Objects: references/objects/{id}.jpg
    obj_dir = refs_dir / "objects"
    if obj_dir.exists():
        for f in obj_dir.iterdir():
            if f.is_file() and f.suffix in (".jpg", ".png"):
                existing_refs.add(("objects", f.stem))

    # Find missing: in episodes but not in references
    missing = {"characters": {}, "locations": {}, "objects": {}}
    has_missing = False

    for asset_type in ["characters", "locations", "objects"]:
        # Map characters -> people for ref check
        ref_type = asset_type
        for aid, info in discovered[asset_type].items():
            if (ref_type, aid) not in existing_refs:
                missing[asset_type][aid] = info
                has_missing = True

    if not has_missing:
        return "All discovered assets have reference sheets."

    # Format output
    lines = ["Missing asset references (not in references/):", ""]

    for asset_type in ["objects", "locations", "characters"]:
        items = missing[asset_type]
        if not items:
            continue

        lines.append(f"  {asset_type}:")
        for aid, info in sorted(items.items()):
            # Format episode/scene info
            ep_parts = []
            for ep_num, scenes in sorted(info["episodes"].items()):
                if scenes:
                    scene_str = _format_scene_list(scenes)
                    ep_parts.append(f"ep{ep_num} scenes {scene_str}")
                else:
                    ep_parts.append(f"ep{ep_num}")
            ep_str = ", ".join(ep_parts)

            # Truncate description
            desc = info["description"]
            if len(desc) > 80:
                desc = desc[:77] + "..."

            in_plan = " [in plan]" if aid in plan_ids else ""
            lines.append(f'    {aid:<25} ({ep_str}) "{desc}"{in_plan}')
        lines.append("")

    lines.append("To create a sheet: add-asset <type> <id>")
    lines.append('  e.g.: add-asset objects sandwolves')

    return "\n".join(lines)


def add_asset(story_path: Path, asset_type: str, asset_id: str,
              description: str = None, dep_ids: list = None) -> dict:
    """Add a new asset reference sheet.

    1. Check if already in plan — warn if exists
    2. Gather description from episodes if not provided
    3. Call Claude to create NanoBanana prompt
    4. Add to assets_plan.json
    5. Generate sheet via _generate_from_prompt()
    6. Auto-select + copy to references

    Returns:
        {"status": "success"/"error", ...}
    """
    if asset_type not in ASSET_TYPES:
        return {"status": "error",
                "error": f"Invalid type '{asset_type}'. "
                         f"Must be one of: {', '.join(ASSET_TYPES)}"}

    stage_dir = story_path / "S5-assets"
    plan_file = stage_dir / "assets_plan.json"

    if not plan_file.exists():
        return {"status": "error",
                "error": "No assets plan. Run S5-assets first."}

    plan = json.loads(plan_file.read_text())

    # Check if already in plan
    existing_type, existing_info = _find_asset(plan, asset_id)
    if existing_type:
        # Check if it already has a reference sheet
        asset_dir = stage_dir / existing_type / asset_id
        iters = _list_asset_iterations(asset_dir)
        if iters:
            return {"status": "error",
                    "error": f"Asset '{asset_id}' already exists in plan "
                             f"as {existing_type} with {len(iters)} iteration(s). "
                             f"Use 'regenerate S5-assets {asset_id}' instead."}
        # In plan but no iterations — we'll generate
        print(f"Asset '{asset_id}' is in plan but has no iterations. Generating...")
        asset_type = existing_type

    # Gather description from episodes if not provided
    if not description:
        description = _find_asset_description_in_episodes(
            story_path, asset_type, asset_id)

    if not description:
        return {"status": "error",
                "error": f"No description provided and could not find "
                         f"'{asset_id}' in episode 04-assets files. "
                         f"Provide a description as the third argument."}

    # Load session
    session_file = story_path / "session.txt"
    if not session_file.exists():
        return {"status": "error", "error": "session.txt not found"}
    session_id = session_file.read_text().strip()

    # Resolve paths for Claude
    bible_path = _find_bible_path(story_path)
    style_path = _find_style_path(story_path)

    # Build prompt for Claude to create NanoBanana prompt
    asset_dir_rel = f"S5-assets/{asset_type}/{asset_id}"
    prompt_output_rel = f"{asset_type}/{asset_id}/prompt.txt"

    claude_prompt = (
        f"Read the story bible at {bible_path} and visual style at {style_path}.\n"
        f"Also study the style board at references/style/style_board.jpg.\n\n"
        f"Create a NanoBanana reference sheet prompt for this {asset_type[:-1]}:\n\n"
        f"ID: {asset_id}\n"
        f"Description: {description}\n\n"
    )

    if asset_type == "characters":
        claude_prompt += (
            "Write the NanoBanana prompt to: "
            f"S5-assets/{prompt_output_rel}\n\n"
            "Also write a character description to: "
            f"S5-assets/{asset_type}/{asset_id}/description.md\n"
            "Include: Physical Appearance, Clothing, Voice, "
            "Personality Notes, Continuity Notes.\n\n"
        )
    else:
        claude_prompt += (
            "Write the NanoBanana prompt to: "
            f"S5-assets/{prompt_output_rel}\n\n"
        )

    claude_prompt += (
        "Follow the same style and format as existing prompts in S5-assets/. "
        "Look at a few existing prompts for reference.\n"
        "The prompt should create a reference sheet showing the asset from "
        "multiple angles/states in the project's visual style.\n"
        "Write ONLY the prompt file(s). Do NOT modify assets_plan.json."
    )

    print(f"\nAdding asset: {asset_type}/{asset_id}")
    print(f"Description: {description[:100]}{'...' if len(description) > 100 else ''}")
    print(f"Session: {session_id}")
    print("Calling Claude to create NanoBanana prompt...")

    try:
        result = subprocess.run(
            ["claude", "--resume", session_id,
             "--allowedTools", "Read,Write,Glob,Grep",
             "-p", claude_prompt],
            capture_output=True, text=True,
            cwd=str(story_path), timeout=300,
            env=claude_env()
        )
        if result.returncode != 0:
            error_detail = result.stderr or result.stdout
            return {"status": "error",
                    "error": f"Claude returned code {result.returncode}: "
                             f"{error_detail[:500]}"}
    except subprocess.TimeoutExpired:
        return {"status": "error", "error": "Claude timed out"}

    # Verify prompt was created
    prompt_file = stage_dir / prompt_output_rel
    if not prompt_file.exists():
        return {"status": "error",
                "error": f"Claude did not create prompt file at "
                         f"S5-assets/{prompt_output_rel}"}

    # Add to assets_plan.json if not already there
    if not existing_type:
        new_entry = {
            "id": asset_id,
            "name": asset_id.replace("_", " ").title(),
            "depends_on": dep_ids or [],
            "prompt_file": prompt_output_rel,
        }
        if asset_type == "characters":
            desc_file_rel = f"{asset_type}/{asset_id}/description.md"
            new_entry["description_file"] = desc_file_rel

        if asset_type not in plan:
            plan[asset_type] = []
        plan[asset_type].append(new_entry)
        plan_file.write_text(json.dumps(plan, indent=2) + "\n")
        print(f"Added to assets_plan.json")
    else:
        print(f"Already in assets_plan.json")

    # Resolve references
    style_refs = _get_style_references(story_path)
    asset_info = _find_asset(plan, asset_id)[1]
    dep_refs = _resolve_dependency_refs(stage_dir, plan, asset_info) if asset_info else []
    all_refs = (style_refs or []) + dep_refs
    ref_images = all_refs if all_refs else None

    # Read description file for characters
    desc_file = None
    if asset_type == "characters":
        desc_path = stage_dir / asset_type / asset_id / "description.md"
        if desc_path.exists():
            desc_file = desc_path

    # Generate sheet
    print(f"Generating reference sheet...")
    asset_dir = stage_dir / asset_type / asset_id

    result_path = _generate_asset_sheet(
        asset_dir=asset_dir,
        prompt_file=prompt_file,
        desc_file=desc_file,
        ref_images=ref_images,
        asset_type=asset_type,
        asset_id=asset_id,
        story_path=story_path,
    )

    if not result_path:
        return {"status": "error",
                "error": "NanoBanana generation failed"}

    iteration = result_path.parent.name

    # Update selections
    selections = load_selections(stage_dir)
    if asset_type not in selections:
        selections[asset_type] = {}
    selections[asset_type][asset_id] = iteration
    save_selections(stage_dir, selections)

    # Copy to references
    _copy_to_references(story_path, asset_type, asset_id, result_path)

    ref_path = _get_reference_path(story_path, asset_type, asset_id)
    print(f"\nDone: {asset_type}/{asset_id}/{iteration}/sheet.jpg")
    print(f"  -> {ref_path.relative_to(story_path)}")

    return {
        "status": "success",
        "asset_type": asset_type,
        "asset_id": asset_id,
        "iteration": iteration,
    }


def _format_scene_list(scenes: list) -> str:
    """Format scene numbers into a compact string (e.g., '1-5,8,10-13')."""
    if not scenes:
        return ""
    nums = sorted(int(s) for s in scenes)
    ranges = []
    start = nums[0]
    end = nums[0]
    for n in nums[1:]:
        if n == end + 1:
            end = n
        else:
            ranges.append(f"{start}" if start == end else f"{start}-{end}")
            start = end = n
    ranges.append(f"{start}" if start == end else f"{start}-{end}")
    return ",".join(ranges)


def _find_asset_description_in_episodes(story_path: Path, asset_type: str,
                                         asset_id: str) -> str | None:
    """Search episode 04-assets files for an asset's description."""
    for ep_dir in sorted(story_path.iterdir()):
        if not ep_dir.is_dir() or not ep_dir.name.isdigit() or len(ep_dir.name) != 3:
            continue

        assets_stage = ep_dir / "04-assets"
        if not assets_stage.exists():
            continue

        iters = sorted(
            [d for d in assets_stage.iterdir() if d.is_dir() and d.name.isdigit()],
            key=lambda d: d.name
        )
        if not iters:
            continue

        assets_file = iters[-1] / "assets.json"
        if not assets_file.exists():
            continue

        try:
            data = json.loads(assets_file.read_text())
        except (json.JSONDecodeError, OSError):
            continue

        for asset in data.get(asset_type, []):
            if asset.get("id") == asset_id:
                return asset.get("description", asset.get("appearance"))

    return None


def _find_bible_path(story_path: Path) -> str | None:
    """Find the story bible path from pipeline state."""
    from lib.story_pipeline_state import load_story_pipeline
    pipeline = load_story_pipeline(story_path)
    bible_iter = pipeline.get("S2-bible")
    if bible_iter:
        return f"S2-bible/{bible_iter}/story_bible.json"
    return None


def _find_style_path(story_path: Path) -> str | None:
    """Find the style path from pipeline state."""
    from lib.story_pipeline_state import load_story_pipeline
    pipeline = load_story_pipeline(story_path)
    style_iter = pipeline.get("S3-style")
    if style_iter:
        return f"S3-style/{style_iter}/style.json"
    return None


def _get_reference_path(story_path: Path, asset_type: str,
                        asset_id: str) -> Path:
    """Get the expected reference path for an asset."""
    refs_path = story_path / "references"
    if asset_type == "characters":
        return refs_path / "people" / asset_id / "sheet.jpg"
    elif asset_type == "locations":
        return refs_path / "locations" / f"{asset_id}.jpg"
    else:
        return refs_path / "objects" / f"{asset_id}.jpg"


def _copy_to_references(story_path: Path, asset_type: str,
                        asset_id: str, sheet_path: Path) -> None:
    """Copy generated sheet + description to references directory."""
    refs_path = story_path / "references"
    iter_dir = sheet_path.parent

    if asset_type == "characters":
        ref_dir = refs_path / "people" / asset_id
        ensure_dir(ref_dir)
        shutil.copy2(sheet_path, ref_dir / "sheet.jpg")
        # Also copy as portrait.jpg for backward compatibility
        shutil.copy2(sheet_path, ref_dir / "portrait.jpg")

        # Copy description.md from iteration dir
        desc_file = iter_dir / "description.md"
        if desc_file.exists():
            shutil.copy2(desc_file, ref_dir / "description.md")

    elif asset_type == "locations":
        ref_dir = refs_path / "locations"
        ensure_dir(ref_dir)
        shutil.copy2(sheet_path, ref_dir / f"{asset_id}.jpg")

    elif asset_type == "objects":
        ref_dir = refs_path / "objects"
        ensure_dir(ref_dir)
        shutil.copy2(sheet_path, ref_dir / f"{asset_id}.jpg")
