#!/home/videobot/dev/venv/bin/python3
"""
Story Pipeline: Multi-Episode Literary Adaptation

Sits above orchestrator_v5.py. Handles story analysis, character design,
title artwork, episode splitting. Before each episode it prepares context,
after each episode it feeds results back for continuity.

Story-level stages:
    S1: Ingest (story text + references)
    S2: Story Bible (characters, locations, themes)
    S3: Style + Style Board (visual direction, NanoBanana style board)
    S4: Title Artwork (series title, NanoBanana with style board reference)
    S5: Asset Reference Sheets (per-asset iterations: characters, locations, objects)
    S6: Episode Breakdown (split story into episodes)

Per-episode stages:
    S7: Prepare Episode (create brief.txt, title card)
    S8: Review Episode (post-assembly continuity update)

Commands:
    init URL_OR_FILE [--style "..."]      Initialize story project
    init-sketch-channel [--style "..."]   Initialize sketch channel project
    run                                    Generate next stage / show pending
    run --accept ["feedback"]              Accept current and advance
    run --refine "feedback"                Refine latest iteration
    run --regenerate ["feedback"]           Fresh new iteration
    chat STAGE [TARGET] "prompt"           Refine NanoBanana image via chat
    refine STAGE TARGET "feedback"         Refine asset with new prompt
    regenerate STAGE TARGET ["feedback"]   Regenerate asset fresh
    select STAGE [TARGET] ITERATION        Select iteration
    deps ASSET_ID [DEP_IDS...]             Show/set asset dependencies
    add-asset --list                       List missing asset references
    add-asset TYPE ID ["description"]      Generate new reference sheet
    status                                 Show pipeline state
    prepare-episode NNN                    Prepare episode for orchestrator
    review-episode NNN                     Post-episode continuity review
    new-sketch "premise"                   Create new sketch (sketch channels only)
"""

import argparse
import json
import os
import subprocess
import sys
import uuid
from pathlib import Path
from dotenv import load_dotenv

# Load .env from repository root
load_dotenv(Path(__file__).parent.parent.parent / ".env")

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

from lib import claude_env
from lib.path_utils_v3 import ensure_dir, get_next_iteration
from lib.iteration import list_iterations, read_metadata
from lib.story_pipeline_state import (
    STORY_STAGES,
    load_story_pipeline,
    save_story_pipeline,
    find_next_story_stage,
    get_project_type,
    get_episode_state,
    set_episode_state,
    get_episodes_from_breakdown,
    format_story_status,
)

from story_stages import (
    stage_s1_ingest,
    stage_s2_bible,
    stage_s3_style,
    stage_s4_title,
    stage_s5_assets,
    stage_s5_episodes as stage_s6_episodes,
    stage_s6_prepare as stage_s7_prepare,
    stage_s7_review as stage_s8_review,
)

# Dependencies for each story stage
STORY_STAGE_DEPENDENCIES = {
    "S1-ingest":      [],
    "S2-bible":       ["S1-ingest"],
    "S3-style":       ["S2-bible"],
    "S4-title":       ["S2-bible", "S3-style"],
    "S5-assets":      ["S2-bible", "S3-style"],
    "S6-episodes":    ["S2-bible", "S3-style"],
}


def resolve_story_inputs(pipeline: dict, stage: str) -> dict:
    """Get parent stage iterations for a given story stage."""
    deps = STORY_STAGE_DEPENDENCIES.get(stage, [])
    inputs = {}
    for parent_stage in deps:
        parent_iter = pipeline.get(parent_stage)
        if parent_iter is None:
            raise ValueError(
                f"Cannot run {stage}: parent stage {parent_stage} "
                "has no selected iteration"
            )
        inputs[parent_stage] = parent_iter
    return inputs


# ============================================================
# Commands
# ============================================================


def cmd_init(story_path: Path, args) -> int:
    """Initialize story project structure + ingest source."""
    source = args.source
    style = getattr(args, 'style', None)

    print(f"Initializing story project: {story_path.name}")
    print(f"Path: {story_path}")

    # Create project directory
    ensure_dir(story_path)

    # Create session
    session_file = story_path / "session.txt"
    if not session_file.exists():
        session_id = str(uuid.uuid4())
        session_file.write_text(session_id)
        print(f"Session: {session_id}")
    else:
        session_id = session_file.read_text().strip()
        print(f"Session exists: {session_id}")

    # Create references directory structure
    ensure_dir(story_path / "references" / "people")
    ensure_dir(story_path / "references" / "locations")
    ensure_dir(story_path / "references" / "objects")
    ensure_dir(story_path / "references" / "style")

    # Create story stage directories
    for stage in STORY_STAGES:
        ensure_dir(story_path / stage)

    # Initialize story pipeline
    pipeline = load_story_pipeline(story_path)

    # Save style preference if provided
    if style:
        style_file = story_path / "style_preference.txt"
        style_file.write_text(style)
        print(f"Style preference: {style}")

    # Create CLAUDE.md from story-adaptation template
    claude_md = story_path / "CLAUDE.md"
    if not claude_md.exists():
        template_path = (Path(__file__).parent.parent
                         / "templates" / "story-adaptation" / "CLAUDE.md.template")
        if template_path.exists():
            content = template_path.read_text()
            content = content.replace("{story_name}", story_path.name)
            claude_md.write_text(content)
            print("Created: CLAUDE.md")
        else:
            print(f"Warning: CLAUDE.md template not found at {template_path}")

    # Initialize Claude session
    story_name = story_path.name
    print(f"\nCreating Claude session...")
    init_prompt = (
        f'You are starting a multi-episode microdrama adaptation project: '
        f'"{story_name}"\n\n'
        f'This is a story-level pipeline that will:\n'
        f'1. Analyze a short story\n'
        f'2. Design visual style and characters\n'
        f'3. Break into episodes\n'
        f'4. Prepare each episode for the per-episode video pipeline\n\n'
        f'Working directory: {story_path}\n\n'
        f'Acknowledge this session initialization briefly.'
    )

    try:
        result = subprocess.run(
            ["claude", "--session-id", session_id,
             "--allowedTools", "Read,Glob",
             "-p", init_prompt],
            capture_output=True, text=True,
            cwd=str(story_path), timeout=120,
            env=claude_env()
        )
        if result.returncode == 0:
            print("Claude session created")
        else:
            print(f"Warning: session init returned code {result.returncode}")
    except (subprocess.TimeoutExpired, Exception) as e:
        print(f"Warning: session init issue: {e}")

    # Run S1: Ingest
    print(f"\nIngesting source: {source}")
    result = stage_s1_ingest.execute(story_path, source)

    if result["status"] != "success":
        print(f"Error: {result.get('error')}")
        return 1

    # Auto-select S1
    pipeline["S1-ingest"] = result["iteration"]
    save_story_pipeline(story_path, pipeline)
    print(f"S1-ingest: {result['iteration']} ✓")

    print(f"\nProject initialized!")
    print(f"Next: story_pipeline.py {story_path} run")

    return 0


def cmd_init_sketch_channel(story_path: Path, args) -> int:
    """Initialize sketch channel project structure."""
    style = getattr(args, 'style', None)
    channel_name = story_path.name

    print(f"Initializing sketch channel: {channel_name}")
    print(f"Path: {story_path}")

    # Create project directory
    ensure_dir(story_path)

    # Create session
    session_file = story_path / "session.txt"
    if not session_file.exists():
        session_id = str(uuid.uuid4())
        session_file.write_text(session_id)
        print(f"Session: {session_id}")
    else:
        session_id = session_file.read_text().strip()
        print(f"Session exists: {session_id}")

    # Create references directory structure
    ensure_dir(story_path / "references" / "people")
    ensure_dir(story_path / "references" / "locations")
    ensure_dir(story_path / "references" / "objects")
    ensure_dir(story_path / "references" / "style")

    # Create story stage directories (skip S6-episodes for sketch channels)
    for stage in STORY_STAGES:
        if stage == "S6-episodes":
            continue
        ensure_dir(story_path / stage)

    # Initialize story pipeline with sketch_channel type
    pipeline = load_story_pipeline(story_path, project_type="sketch_channel")
    # Mark S6-episodes as N/A for sketch channels
    pipeline["S6-episodes"] = "N/A"
    save_story_pipeline(story_path, pipeline)

    # Save style preference if provided
    if style:
        style_file = story_path / "style_preference.txt"
        style_file.write_text(style)
        print(f"Style preference: {style}")

    # Create CLAUDE.md from sketch-channel template
    claude_md = story_path / "CLAUDE.md"
    if not claude_md.exists():
        template_path = (Path(__file__).parent.parent
                         / "templates" / "sketch-channel" / "CLAUDE.md.template")
        if template_path.exists():
            content = template_path.read_text()
            content = content.replace("{channel_name}", channel_name)
            claude_md.write_text(content)
            print("Created: CLAUDE.md")
        else:
            print(f"Warning: CLAUDE.md template not found at {template_path}")

    # Create channel_concept.txt placeholder
    concept_file = story_path / "channel_concept.txt"
    if not concept_file.exists():
        concept_file.write_text(
            f"# {channel_name}\n\n"
            f"Recurring-character comedy sketch channel.\n\n"
            f"## Concept\n\n"
            f"[Describe the channel premise, tone, and recurring characters here]\n"
        )
        print("Created: channel_concept.txt (edit with your channel concept)")

    # Initialize Claude session
    print(f"\nCreating Claude session...")
    init_prompt = (
        f'You are starting a recurring-character comedy sketch channel project: '
        f'"{channel_name}"\n\n'
        f'This is a sketch channel pipeline that will:\n'
        f'1. Define recurring characters and comedy style\n'
        f'2. Design visual style and character references\n'
        f'3. Prepare standalone comedy sketches for video production\n'
        f'4. Track character behavior and running gags across sketches\n\n'
        f'Key difference from story adaptation: each sketch is standalone — '
        f'same characters but no continuous storyline.\n\n'
        f'Working directory: {story_path}\n\n'
        f'Acknowledge this session initialization briefly.'
    )

    try:
        result = subprocess.run(
            ["claude", "--session-id", session_id,
             "--allowedTools", "Read,Glob",
             "-p", init_prompt],
            capture_output=True, text=True,
            cwd=str(story_path), timeout=120,
            env=claude_env()
        )
        if result.returncode == 0:
            print("Claude session created")
        else:
            print(f"Warning: session init returned code {result.returncode}")
    except (subprocess.TimeoutExpired, Exception) as e:
        print(f"Warning: session init issue: {e}")

    print(f"\nSketch channel initialized!")
    print(f"\nNext steps:")
    print(f"  1. Edit channel_concept.txt with your channel premise")
    print(f"  2. Run S1-S5 stages (S6-episodes is skipped for sketch channels):")
    print(f"     story_pipeline.py {story_path} run")
    print(f"  3. Use 'new-sketch' to create sketches:")
    print(f"     story_pipeline.py {story_path} new-sketch \"Drill Sergeant tries to read a map\"")

    return 0


def cmd_new_sketch(story_path: Path, args) -> int:
    """Create a new sketch episode from a premise."""
    premise = args.premise
    pipeline = load_story_pipeline(story_path)

    if get_project_type(pipeline) != "sketch_channel":
        print("Error: new-sketch is only for sketch_channel projects.")
        print("This project is a story_adaptation. Use prepare-episode instead.")
        return 1

    # Find next episode number
    existing = sorted([
        d.name for d in story_path.iterdir()
        if d.is_dir() and d.name.isdigit()
    ])
    if existing:
        next_num = int(existing[-1]) + 1
    else:
        next_num = 1
    episode_num = f"{next_num:03d}"

    episode_path = story_path / episode_num
    ensure_dir(episode_path)

    # Create sketch script from premise
    script_file = episode_path / "sketch_script.md"
    script_file.write_text(f"# {premise}\n\n{premise}\n")
    print(f"Created: {episode_num}/sketch_script.md")

    # Also copy to 01-story/001/ for orchestrator
    story_dir = episode_path / "01-story" / "001"
    ensure_dir(story_dir)
    import shutil
    shutil.copy(script_file, story_dir / "story.md")

    print(f"\nSketch {episode_num} created: \"{premise}\"")
    print(f"\nNext steps:")
    print(f"  1. (Optional) Flesh out {episode_num}/sketch_script.md with full dialogue")
    print(f"  2. Prepare the sketch brief:")
    print(f"     story_pipeline.py {story_path} prepare-episode {episode_num}")

    return 0


def cmd_status(story_path: Path, args) -> int:
    """Show story pipeline status."""
    pipeline = load_story_pipeline(story_path)
    print(format_story_status(story_path, pipeline))

    is_sketch = get_project_type(pipeline) == "sketch_channel"
    label = "Sketch" if is_sketch else "Episode"

    # Show episode detail if episodes exist
    episodes = get_episodes_from_breakdown(story_path, pipeline)
    if episodes:
        print(f"\n{label} Details:")
        for ep in episodes:
            num = f"{ep['number']:03d}"
            ep_state = get_episode_state(pipeline, num)
            ep_path = story_path / num
            has_dir = ep_path.exists()
            has_assembly = False
            if has_dir:
                for asm_name in ["13-assembly", "12-assembly"]:
                    if (ep_path / asm_name).exists():
                        has_assembly = True
                        break

            status_parts = []
            if ep_state.get("prepared"):
                status_parts.append("prepared")
            if has_assembly:
                status_parts.append("assembled")
            if ep_state.get("reviewed"):
                status_parts.append("reviewed")

            status_str = (", ".join(status_parts)
                          if status_parts else "not started")
            grid = ep.get("grid_layout", "?")
            dur = ep.get("clip_duration", ep.get("target_duration", "?"))
            print(f"  {num}: \"{ep.get('title', '')}\" "
                  f"({grid}, {dur}) -- {status_str}")

    return 0


def cmd_run(story_path: Path, args) -> int:
    """Run next stage or handle accept/refine/regenerate."""
    pipeline = load_story_pipeline(story_path)
    mode = _get_run_mode(args)

    # Find current stage
    next_stage = find_next_story_stage(pipeline)

    if next_stage is None:
        is_sketch = get_project_type(pipeline) == "sketch_channel"
        if is_sketch:
            print("All channel setup stages complete!")
            print("\nNext steps:")
            print(f"  story_pipeline.py {story_path} new-sketch \"Your sketch premise\"")
            print(f"  story_pipeline.py {story_path} prepare-episode NNN")
        else:
            print("All story stages complete!")
            print("\nNext steps:")
            episodes = get_episodes_from_breakdown(story_path, pipeline)
            for ep in episodes:
                num = f"{ep['number']:03d}"
                ep_state = get_episode_state(pipeline, num)
                if not ep_state.get("prepared"):
                    print(f"  story_pipeline.py {story_path} "
                          f"prepare-episode {num}")
                    break
        return 0

    stage = next_stage

    # --- S5-ASSETS special handling ---
    if stage == "S5-assets":
        return _run_s5_assets(story_path, pipeline, mode, args)

    # --- ACCEPT ---
    if mode == "accept":
        latest = _get_latest_iteration(story_path, stage)
        if not latest:
            print(f"Nothing to accept in {stage} — generating instead...")
            # Fall through to generate
        else:
            pipeline[stage] = latest
            save_story_pipeline(story_path, pipeline)
            feedback = args.accept
            print(f"Accepted: {stage}/{latest}")
            if feedback:
                print(f"Feedback noted: \"{feedback}\"")

            # Show next step
            next_next = find_next_story_stage(pipeline)
            if next_next:
                print(f"\nNext: {next_next}")
                print(f"  story_pipeline.py {story_path} run")
            else:
                print(f"\nAll story stages complete!")
            return 0

    # --- REFINE ---
    if mode == "refine":
        latest = _get_latest_iteration(story_path, stage)
        if not latest:
            print(f"Nothing to refine in {stage}")
            return 1

        feedback = args.refine
        print(f"Refining: {stage}/{latest}")
        print(f"Feedback: {feedback}")

        result = _run_stage(story_path, stage, pipeline,
                            feedback=feedback, refined_from=latest)

        if result["status"] != "success":
            print(f"\nError: {result.get('error')}")
            return 1

        print(f"\nGenerated: {stage}/{result['iteration']} "
              f"(refined from {latest})")
        if result.get("user_message"):
            print(f"\n{result['user_message']}")
        return 0

    # --- REGENERATE ---
    if mode == "regenerate":
        feedback = args.regenerate or None
        print(f"Regenerating: {stage}")

        result = _run_stage(story_path, stage, pipeline, feedback=feedback)

        if result["status"] != "success":
            print(f"\nError: {result.get('error')}")
            return 1

        print(f"\nGenerated: {stage}/{result['iteration']}")
        if result.get("user_message"):
            print(f"\n{result['user_message']}")
        return 0

    # --- GENERATE (default) ---
    # Check for pending (has iteration but not selected)
    iters = list_iterations(story_path / stage)
    real_iters = [i for i in iters
                  if (story_path / stage / i / "metadata.json").exists()]
    if real_iters and pipeline.get(stage) is None:
        latest = real_iters[-1]
        print(f"Pending: {stage}/{latest}")
        _show_iteration_status(story_path, stage, latest)
        return 0

    # Generate
    print(f"Generating: {stage}")
    result = _run_stage(story_path, stage, pipeline)

    if result["status"] != "success":
        print(f"\nError: {result.get('error')}")
        return 1

    iteration = result["iteration"]
    print(f"\nGenerated: {stage}/{iteration}")

    if result.get("user_message"):
        print(f"\n{result['user_message']}")

    print()
    print("Actions:")
    print(f"  run --accept            Accept and advance")
    print(f"  run --refine \"feedback\"  Refine with feedback")
    print(f"  run --regenerate        Generate fresh alternative")

    return 0


def _run_s5_assets(story_path: Path, pipeline: dict,
                   mode: str, args) -> int:
    """Handle S5-assets stage within cmd_run."""

    # --- ACCEPT ---
    if mode == "accept":
        success, msg = stage_s5_assets.accept_all(story_path)
        print(msg)
        if success:
            pipeline["S5-assets"] = "complete"
            save_story_pipeline(story_path, pipeline)
            next_next = find_next_story_stage(pipeline)
            if next_next:
                print(f"\nNext: {next_next}")
                print(f"  story_pipeline.py {story_path} run")
            else:
                print(f"\nAll story stages complete!")
        return 0 if success else 1

    # --- REFINE / REGENERATE at stage level ---
    if mode in ("refine", "regenerate"):
        print(f"S5-assets uses per-asset operations:")
        print(f"  refine S5-assets <asset_id> \"feedback\"")
        print(f"  regenerate S5-assets <asset_id> [\"feedback\"]")
        return 1

    # --- GENERATE (default) ---
    plan_file = story_path / "S5-assets" / "assets_plan.json"

    if plan_file.exists():
        # Plan exists — show status
        print(stage_s5_assets.get_asset_status(story_path))
        print()
        print("Actions:")
        print(f"  run --accept                              "
              f"Accept all and advance")
        print(f"  refine S5-assets <asset_id> \"feedback\"    "
              f"Refine specific asset")
        print(f"  regenerate S5-assets <asset_id>           "
              f"Regenerate specific asset")
        print(f"  chat S5-assets <asset_id> \"prompt\"        "
              f"Chat-refine via NanoBanana")
        print(f"  select S5-assets <asset_id> <iteration>   "
              f"Select iteration")
        return 0

    # No plan yet — generate everything
    result = _run_stage(story_path, "S5-assets", pipeline)

    if result["status"] != "success":
        print(f"\nError: {result.get('error')}")
        return 1

    if result.get("user_message"):
        print(f"\n{result['user_message']}")

    print()
    print(stage_s5_assets.get_asset_status(story_path))
    print()
    print("Actions:")
    print(f"  run --accept                              "
          f"Accept all and advance")
    print(f"  refine S5-assets <asset_id> \"feedback\"    "
          f"Refine specific asset")
    print(f"  regenerate S5-assets <asset_id>           "
          f"Regenerate specific asset")
    print(f"  chat S5-assets <asset_id> \"prompt\"        "
          f"Chat-refine via NanoBanana")
    print(f"  select S5-assets <asset_id> <iteration>   "
          f"Select iteration")

    return 0


def cmd_chat(story_path: Path, args) -> int:
    """Refine a NanoBanana image via chat continuation."""
    stage = args.stage
    prompt_text = args.prompt
    target = getattr(args, 'target', None)

    pipeline = load_story_pipeline(story_path)

    if stage == "S3-style":
        return _chat_single_image(
            story_path, pipeline, "S3-style",
            "style_board.jpg", "chat_history.json",
            "references/style/style_board.jpg",
            prompt_text)

    elif stage == "S4-title":
        return _chat_single_image(
            story_path, pipeline, "S4-title",
            "title_artwork.jpg", "chat_history.json",
            "references/style/title_artwork.jpg",
            prompt_text)

    elif stage == "S5-assets":
        if not target:
            print("S5-assets requires a target asset ID.")
            print('Usage: chat S5-assets <asset_id> "prompt"')
            # List available assets
            plan_file = story_path / "S5-assets" / "assets_plan.json"
            if plan_file.exists():
                plan = json.loads(plan_file.read_text())
                ids = stage_s5_assets._list_all_asset_ids(plan)
                if ids:
                    print(f"Available: {', '.join(ids)}")
            return 1

        result = stage_s5_assets.execute_single(
            story_path, target, "chat", feedback=prompt_text)

        if result["status"] != "success":
            print(f"Error: {result.get('error')}")
            return 1

        print(f"\nChat-refined: {result['asset']}/{result['iteration']}")
        print(f"\nActions:")
        print(f'  chat S5-assets {target} "more refinement"')
        print(f"  select S5-assets {target} {result['iteration']}")
        return 0

    else:
        print(f"Chat not supported for stage: {stage}")
        print("Supported: S3-style, S4-title, S5-assets")
        return 1


def _chat_single_image(story_path: Path, pipeline: dict,
                       stage: str, image_name: str,
                       history_name: str, ref_copy_rel: str,
                       prompt_text: str) -> int:
    """Chat-refine a single-image stage (S3-style or S4-title)."""
    import shutil
    from lib import save_metadata
    from datetime import datetime, timezone

    current_iter = (pipeline.get(stage)
                    or _get_latest_iteration(story_path, stage))
    if not current_iter:
        print(f"No {stage} iteration to refine. Run {stage} first.")
        return 1

    src_dir = story_path / stage / current_iter
    chat_history = src_dir / history_name

    if not chat_history.exists():
        print(f"No chat history found at: {chat_history}")
        return 1

    # Create new iteration
    stage_dir = story_path / stage
    new_iter = get_next_iteration(stage_dir)
    new_iter_dir = stage_dir / new_iter
    ensure_dir(new_iter_dir)

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

    new_chat_history = new_iter_dir / history_name
    new_output = new_iter_dir / image_name
    ref_copy_to = story_path / ref_copy_rel

    # Generate refined image
    print(f"Refining {image_name} in {stage}/{new_iter}...")
    print(f"  Chat: {prompt_text}")

    try:
        repo_root = Path(__file__).parent.parent.parent
        sys.path.insert(0, str(repo_root / "scripts"))
        from nanobanana_text2img import generate_image_chat

        result_path = generate_image_chat(
            prompt=prompt_text,
            output_path=str(new_output),
            history_path=str(new_chat_history),
            aspect_ratio="9:16",
            image_size="4K",
            is_refinement=True,
        )
        print(f"  Generated: {Path(result_path).name}")

        # Copy to references
        if ref_copy_to:
            ensure_dir(ref_copy_to.parent)
            shutil.copy2(result_path, ref_copy_to)
            print(f"  -> {ref_copy_to.relative_to(story_path)}")

        # Save metadata
        metadata = {
            "stage": stage,
            "iteration": new_iter,
            "created_at": datetime.now(timezone.utc).isoformat(),
            "refined_from": current_iter,
            "chat_refinement": prompt_text,
        }
        save_metadata(new_iter_dir / "metadata.json", metadata)

        print(f"\nRefined: {stage}/{new_iter}/{image_name}")
        print(f"\nActions:")
        print(f'  chat {stage} "more refinement"')
        print(f"  select {stage} {new_iter}")

        # Auto-select if stage was already selected
        if pipeline.get(stage):
            pipeline[stage] = new_iter
            save_story_pipeline(story_path, pipeline)
            print(f"  Auto-selected: {stage}/{new_iter}")

        return 0

    except Exception as e:
        print(f"Error: NanoBanana refinement failed: {e}")
        shutil.rmtree(new_iter_dir, ignore_errors=True)
        return 1


def cmd_refine(story_path: Path, args) -> int:
    """Refine a specific asset with feedback."""
    stage = args.stage

    if stage == "S5-assets":
        asset_id = args.target_or_feedback
        feedback = args.extra_feedback
        if not feedback:
            print("Usage: refine S5-assets <asset_id> \"feedback\"")
            return 1

        result = stage_s5_assets.execute_single(
            story_path, asset_id, "refine", feedback=feedback)

        if result["status"] != "success":
            print(f"Error: {result.get('error')}")
            return 1

        print(f"\nRefined: {result['asset']}/{result['iteration']}")
        return 0

    else:
        # For other stages, redirect to run --refine
        print(f"Use 'run --refine \"feedback\"' for stage {stage}")
        print(f"Or for S5-assets: refine S5-assets <asset_id> \"feedback\"")
        return 1


def cmd_regenerate(story_path: Path, args) -> int:
    """Regenerate a specific asset."""
    stage = args.stage

    if stage == "S5-assets":
        target = args.target_or_feedback
        feedback = args.extra_feedback
        if not target:
            print("Usage: regenerate S5-assets <asset_id> [\"feedback\"]")
            return 1

        result = stage_s5_assets.execute_single(
            story_path, target, "regenerate", feedback=feedback)

        if result["status"] != "success":
            print(f"Error: {result.get('error')}")
            return 1

        print(f"\nRegenerated: {result['asset']}/{result['iteration']}")
        return 0

    else:
        print(f"Use 'run --regenerate' for stage {stage}")
        print(f"Or for S5-assets: regenerate S5-assets <asset_id>")
        return 1


def cmd_select(story_path: Path, args) -> int:
    """Manually select an iteration for a stage or asset."""
    stage = args.stage
    iteration = args.iteration
    extra = getattr(args, 'extra', None)

    pipeline = load_story_pipeline(story_path)

    if stage == "S5-assets":
        # select S5-assets <asset_id> <iteration>
        asset_id = iteration
        actual_iteration = extra
        if not actual_iteration:
            print("Usage: select S5-assets <asset_id> <iteration>")
            return 1

        success, msg = stage_s5_assets.select_asset(
            story_path, asset_id, actual_iteration)
        print(msg)
        return 0 if success else 1

    # Standard stage selection
    if stage not in STORY_STAGES:
        print(f"Unknown stage: {stage}")
        print(f"Available: {', '.join(STORY_STAGES)}")
        return 1

    iter_path = story_path / stage / iteration
    if not iter_path.exists():
        print(f"Iteration not found: {stage}/{iteration}")
        return 1

    pipeline[stage] = iteration
    save_story_pipeline(story_path, pipeline)
    print(f"Selected: {stage}/{iteration}")
    return 0


def cmd_deps(story_path: Path, args) -> int:
    """Show or set asset dependencies."""
    asset_id = args.asset_id
    dep_ids = args.dep_ids
    clear = args.clear

    plan_file = story_path / "S5-assets" / "assets_plan.json"
    if not plan_file.exists():
        print("No assets plan. Run S5-assets first.")
        return 1

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

    # Find asset
    asset_type = None
    asset_info = None
    for atype in stage_s5_assets.ASSET_TYPES:
        for info in plan.get(atype, []):
            if info["id"] == asset_id:
                asset_type = atype
                asset_info = info
                break
        if asset_type:
            break

    if not asset_type:
        all_ids = []
        for atype in stage_s5_assets.ASSET_TYPES:
            all_ids.extend(i["id"] for i in plan.get(atype, []))
        print(f"Asset '{asset_id}' not found.")
        print(f"Available: {', '.join(all_ids)}")
        return 1

    # Show mode (no dep_ids, no --clear)
    if not dep_ids and not clear:
        current = asset_info.get("depends_on", [])
        if current:
            print(f"{asset_type}/{asset_id} depends on: {', '.join(current)}")
        else:
            print(f"{asset_type}/{asset_id} has no dependencies")
        return 0

    # Clear mode
    if clear:
        asset_info["depends_on"] = []
        plan_file.write_text(json.dumps(plan, indent=2) + "\n")
        print(f"Cleared dependencies for {asset_type}/{asset_id}")
        return 0

    # Set mode — validate all dep_ids exist in plan
    all_ids = set()
    for atype in stage_s5_assets.ASSET_TYPES:
        for info in plan.get(atype, []):
            all_ids.add(info["id"])

    invalid = [d for d in dep_ids if d not in all_ids]
    if invalid:
        print(f"Unknown asset IDs: {', '.join(invalid)}")
        print(f"Available: {', '.join(sorted(all_ids))}")
        return 1

    if asset_id in dep_ids:
        print(f"Asset cannot depend on itself")
        return 1

    asset_info["depends_on"] = dep_ids
    plan_file.write_text(json.dumps(plan, indent=2) + "\n")
    print(f"Set {asset_type}/{asset_id} depends_on: {', '.join(dep_ids)}")
    return 0


def cmd_add_asset(story_path: Path, args) -> int:
    """Add a new asset reference sheet, or list missing assets."""
    if args.list:
        print(stage_s5_assets.list_missing_assets(story_path))
        return 0

    if not args.type or not args.id:
        print("Usage:")
        print("  add-asset --list                          "
              "List missing assets")
        print("  add-asset <type> <id> [\"description\"]     "
              "Generate reference sheet")
        print()
        print("Types: characters, locations, objects")
        return 1

    result = stage_s5_assets.add_asset(
        story_path,
        asset_type=args.type,
        asset_id=args.id,
        description=args.description,
        dep_ids=args.deps if args.deps else None,
    )

    if result["status"] != "success":
        print(f"Error: {result.get('error')}")
        return 1

    print()
    print("Actions:")
    print(f'  chat S5-assets {args.id} "refinement"   '
          f'Chat-refine the sheet')
    print(f'  regenerate S5-assets {args.id}           '
          f'Regenerate from scratch')
    print(f'  add-asset --list                         '
          f'Check remaining missing assets')
    return 0


def cmd_search_corpus(story_path: Path, args) -> int:
    """Search project corpus for relevant reference images."""
    import re
    import shutil
    from datetime import datetime, timezone

    episode = args.episode
    query = args.query

    pipeline = load_story_pipeline(story_path)
    corpus_cfg = pipeline.get("corpus")
    if not corpus_cfg:
        print("Error: No 'corpus' config in story_pipeline.json")
        print("Add: \"corpus\": {\"data_dir\": \"...\", \"images_base\": \"...\", \"top_n\": 15}")
        return 1

    data_dir = corpus_cfg["data_dir"]
    images_base = corpus_cfg.get("images_base")
    top_n = corpus_cfg.get("top_n", 15)

    # If no query provided, use first line of sketch_script.md
    episode_path = story_path / episode
    if not query:
        for candidate in [
            episode_path / "sketch_script.md",
            episode_path / "01-story" / "001" / "story.md",
        ]:
            if candidate.exists():
                first_line = candidate.read_text().strip().split("\n")[0]
                query = first_line.lstrip("# ").strip()
                print(f"Query from script: \"{query}\"")
                break
        if not query:
            print(f"Error: No query provided and no script found in {episode}/")
            return 1

    # Import and search
    from lib.corpus.search import search as corpus_search

    print(f"Searching corpus: \"{query}\" (top {top_n})")
    results = corpus_search(query, data_dir=data_dir, images_base=images_base,
                            top_n=top_n)

    if not results:
        print("No results found.")
        return 1

    # Create refs/ directory and copy images
    refs_dir = episode_path / "refs"
    if refs_dir.exists():
        shutil.rmtree(refs_dir)
    ensure_dir(refs_dir)

    def _sanitize(text: str, max_len: int = 30) -> str:
        text = re.sub(r'[^\w\s-]', '', text.lower())
        text = re.sub(r'[\s]+', '_', text.strip())
        return text[:max_len].rstrip('_')

    for r in results:
        src = Path(r["image_path"])
        if src.exists():
            sanitized = _sanitize(r["title"])
            dest_name = f"{r['rank']:02d}_{sanitized}{src.suffix}"
            shutil.copy(src, refs_dir / dest_name)
            r["local_ref"] = f"refs/{dest_name}"
        else:
            r["local_ref"] = None

    # Write corpus_refs.json
    corpus_refs = {
        "query": query,
        "searched_at": datetime.now(timezone.utc).isoformat(),
        "results": [
            {
                "rank": r["rank"],
                "score": round(r["score"], 4),
                "id": r["id"],
                "title": r["title"],
                "description": r["description"][:300] if r["description"] else "",
                "source_image": r["image_path"],
                "local_ref": r["local_ref"],
            }
            for r in results
        ],
        "scene_assignments": {},
    }
    refs_file = episode_path / "corpus_refs.json"
    refs_file.write_text(json.dumps(corpus_refs, ensure_ascii=False, indent=2))

    # Print results table
    print(f"\nTop {len(results)} results:\n")
    copied = 0
    for r in results:
        print(f"  {r['rank']:2d}. [{r['score']:.3f}] {r['title']}")
        print(f"      Year: {r['year']} | Source: {r['source']}")
        if r["local_ref"]:
            print(f"      Ref: {r['local_ref']}")
            copied += 1
        print()

    print(f"Saved: {refs_file}")
    print(f"Copied: {copied}/{len(results)} images to {refs_dir}/")
    if args.json_output:
        print(f"\nJSON: {refs_file}")

    return 0


def cmd_prepare_episode(story_path: Path, args) -> int:
    """Prepare an episode for the orchestrator pipeline."""
    episode = args.episode
    cli_command = getattr(args, "cli", "claude")
    pipeline = load_story_pipeline(story_path)

    result = stage_s7_prepare.execute(story_path, episode, pipeline,
                                      cli_command=cli_command)

    if result["status"] != "success":
        print(f"Error: {result.get('error')}")
        return 1

    # Update episode state
    set_episode_state(pipeline, episode, "prepared", "001")
    save_story_pipeline(story_path, pipeline)

    if result.get("user_message"):
        print(f"\n{result['user_message']}")

    return 0


def cmd_review_episode(story_path: Path, args) -> int:
    """Review a completed episode and update continuity."""
    episode = args.episode
    pipeline = load_story_pipeline(story_path)

    result = stage_s8_review.execute(story_path, episode, pipeline)

    if result["status"] != "success":
        print(f"Error: {result.get('error')}")
        return 1

    # Update episode state
    set_episode_state(pipeline, episode, "reviewed", "001")
    save_story_pipeline(story_path, pipeline)

    if result.get("user_message"):
        print(f"\n{result['user_message']}")

    return 0


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


def _get_run_mode(args) -> str:
    """Determine run mode from args."""
    if hasattr(args, 'accept') and args.accept is not None:
        return "accept"
    if hasattr(args, 'refine') and args.refine is not None:
        return "refine"
    if hasattr(args, 'regenerate') and args.regenerate is not None:
        return "regenerate"
    return "generate"


def _get_latest_iteration(story_path: Path, stage: str) -> str | None:
    """Get the latest iteration for a story stage."""
    iters = list_iterations(story_path / stage)
    return iters[-1] if iters else None


def _show_iteration_status(story_path: Path, stage: str,
                           iteration: str) -> None:
    """Show status of a pending iteration."""
    iter_dir = story_path / stage / iteration
    user_msg_file = iter_dir / "user_message.txt"

    if user_msg_file.exists():
        print()
        print(user_msg_file.read_text())
    else:
        meta = read_metadata(iter_dir)
        if meta:
            print(f"\n  Stage: {stage}/{iteration}")
            if meta.get("refined_from"):
                print(f"  Refined from: {meta['refined_from']}")

    print()
    print("Actions:")
    print(f"  run --accept            Accept and advance")
    print(f"  run --refine \"feedback\"  Refine with feedback")
    print(f"  run --regenerate        Generate fresh alternative")


def _run_stage(story_path: Path, stage: str, pipeline: dict,
               feedback: str = None, refined_from: str = None) -> dict:
    """Dispatch to the appropriate stage executor."""
    try:
        inputs = resolve_story_inputs(pipeline, stage)
    except ValueError as e:
        return {"status": "error", "error": str(e)}

    if stage == "S1-ingest":
        # For sketch channels, S1-ingest just bootstraps from channel_concept.txt
        if pipeline.get("project_type") == "sketch_channel":
            concept_file = story_path / "channel_concept.txt"
            if not concept_file.exists():
                return {"status": "error",
                        "error": "channel_concept.txt not found. Create it first."}
            # Copy concept as source.txt
            source_file = story_path / "source.txt"
            if not source_file.exists():
                source_file.write_text(concept_file.read_text())
            # Create S1 iteration
            iter_dir = story_path / "S1-ingest" / "001"
            iter_dir.mkdir(parents=True, exist_ok=True)
            import json
            from datetime import datetime, timezone
            (iter_dir / "metadata.json").write_text(json.dumps({
                "stage": "S1-ingest",
                "iteration": "001",
                "created_at": datetime.now(timezone.utc).isoformat(),
                "source": "channel_concept.txt",
                "note": "Bootstrapped from channel concept"
            }, indent=2))
            return {"status": "success", "iteration": "001"}
        return {"status": "error",
                "error": "S1-ingest runs during init. Use 'init' command."}

    elif stage == "S2-bible":
        return stage_s2_bible.execute(
            story_path, inputs,
            feedback=feedback, refined_from=refined_from
        )

    elif stage == "S3-style":
        style_pref = None
        style_file = story_path / "style_preference.txt"
        if style_file.exists():
            style_pref = style_file.read_text().strip()
        return stage_s3_style.execute(
            story_path, inputs,
            feedback=feedback, refined_from=refined_from,
            style_preference=style_pref
        )

    elif stage == "S4-title":
        return stage_s4_title.execute(
            story_path, inputs,
            feedback=feedback, refined_from=refined_from
        )

    elif stage == "S5-assets":
        return stage_s5_assets.execute(
            story_path, inputs,
            feedback=feedback, refined_from=refined_from
        )

    elif stage == "S6-episodes":
        return stage_s6_episodes.execute(
            story_path, inputs,
            feedback=feedback, refined_from=refined_from
        )

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


# ============================================================
# Main
# ============================================================


def main():
    parser = argparse.ArgumentParser(
        description="Story Pipeline: Multi-Episode Literary Adaptation",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Initialize from Gutenberg URL
  story_pipeline.py /mnt/video/microdrama-v3/prospectors-special init \\
    https://www.gutenberg.org/cache/epub/51768/pg51768.txt \\
    --style "retro-futuristic cartoon, Moebius meets Archer"

  # Run through story stages
  story_pipeline.py .../prospectors-special run
  story_pipeline.py .../prospectors-special run --accept
  story_pipeline.py .../prospectors-special run --refine "make it darker"

  # Refine style board via NanoBanana chat
  story_pipeline.py .../prospectors-special chat S3-style "warmer colors"

  # Refine title artwork
  story_pipeline.py .../prospectors-special chat S4-title "remove the kite"

  # Per-asset operations (S5-assets)
  story_pipeline.py .../prospectors-special chat S5-assets tom_morrison "more weathered"
  story_pipeline.py .../prospectors-special refine S5-assets tom_morrison "add scars"
  story_pipeline.py .../prospectors-special regenerate S5-assets scorpion_desert
  story_pipeline.py .../prospectors-special select S5-assets tom_morrison 002

  # Add missing asset reference sheets
  story_pipeline.py .../prospectors-special add-asset --list                   # discover missing
  story_pipeline.py .../prospectors-special add-asset objects sandwolves       # auto-detect description
  story_pipeline.py .../prospectors-special add-asset objects sandwolves "Three lean predators..."

  # Manage asset dependencies
  story_pipeline.py .../prospectors-special deps sandcar_interior              # show deps
  story_pipeline.py .../prospectors-special deps sandcar_interior sandcar telephone  # set deps
  story_pipeline.py .../prospectors-special deps sandcar_interior --clear      # clear deps

  # Prepare episode 1
  story_pipeline.py .../prospectors-special prepare-episode 001

  # After episode is assembled, review it
  story_pipeline.py .../prospectors-special review-episode 001

  # --- Sketch Channel ---

  # Initialize sketch channel
  story_pipeline.py /mnt/video/microdrama-v3/hurry-up-and-wait init-sketch-channel \\
    --style "military-realism"

  # Create a new sketch
  story_pipeline.py .../hurry-up-and-wait new-sketch "Drill Sergeant tries to read a map upside down"

  # Prepare and produce the sketch
  story_pipeline.py .../hurry-up-and-wait prepare-episode 002
  orchestrator_v5.py .../hurry-up-and-wait --episode 002 run

  # Review sketch (updates character notes + running gags)
  story_pipeline.py .../hurry-up-and-wait review-episode 002
""")

    parser.add_argument("story_path", type=Path,
                        help="Path to story project directory")

    subparsers = parser.add_subparsers(dest="command")

    # init
    init_parser = subparsers.add_parser("init",
                                        help="Initialize story project")
    init_parser.add_argument("source",
                             help="URL or file path for story text")
    init_parser.add_argument("--style",
                             help="Style preference")

    # init-sketch-channel
    init_sketch_parser = subparsers.add_parser(
        "init-sketch-channel",
        help="Initialize sketch channel project")
    init_sketch_parser.add_argument("--style",
                                     help="Style preference")

    # run
    run_parser = subparsers.add_parser("run", help="Run next stage")
    run_group = run_parser.add_mutually_exclusive_group()
    run_group.add_argument("--accept", nargs="?", const="", default=None,
                           help="Accept current and advance")
    run_group.add_argument("--refine", help="Refine with feedback")
    run_group.add_argument("--regenerate", nargs="?", const="",
                           default=None,
                           help="Generate fresh alternative")

    # chat
    chat_parser = subparsers.add_parser(
        "chat", help="Refine NanoBanana image via chat")
    chat_parser.add_argument(
        "stage", help="Stage (S3-style, S4-title, S5-assets)")
    chat_parser.add_argument(
        "prompt_or_target",
        help="Refinement prompt, or target ID for S5-assets")
    chat_parser.add_argument(
        "prompt_extra", nargs="?", default=None,
        help="Refinement prompt (when first arg is target ID)")

    # refine (per-asset)
    refine_parser = subparsers.add_parser(
        "refine", help="Refine asset with new prompt")
    refine_parser.add_argument("stage", help="Stage (S5-assets)")
    refine_parser.add_argument(
        "target_or_feedback",
        help="Asset ID for S5-assets")
    refine_parser.add_argument(
        "extra_feedback", nargs="?", default=None,
        help="Feedback text")

    # regenerate (per-asset)
    regen_parser = subparsers.add_parser(
        "regenerate", help="Regenerate asset fresh")
    regen_parser.add_argument("stage", help="Stage (S5-assets)")
    regen_parser.add_argument(
        "target_or_feedback", nargs="?", default=None,
        help="Asset ID for S5-assets")
    regen_parser.add_argument(
        "extra_feedback", nargs="?", default=None,
        help="Optional feedback")

    # status
    subparsers.add_parser("status", help="Show pipeline status")

    # select
    select_parser = subparsers.add_parser(
        "select", help="Select stage iteration")
    select_parser.add_argument("stage",
                               help="Stage name (e.g., S3-style, S5-assets)")
    select_parser.add_argument("iteration",
                               help="Iteration number, or asset ID "
                                    "for S5-assets")
    select_parser.add_argument("extra", nargs="?", default=None,
                               help="Iteration number for S5-assets")

    # add-asset
    add_asset_parser = subparsers.add_parser(
        "add-asset",
        help="Add asset reference sheet (or list missing)")
    add_asset_parser.add_argument(
        "type", nargs="?", default=None,
        choices=["characters", "locations", "objects"],
        help="Asset type")
    add_asset_parser.add_argument(
        "id", nargs="?", default=None,
        help="Asset ID (e.g., sandwolves)")
    add_asset_parser.add_argument(
        "description", nargs="?", default=None,
        help="Optional description (auto-detected from episodes if omitted)")
    add_asset_parser.add_argument(
        "--list", action="store_true",
        help="List assets discovered in episodes but missing references")
    add_asset_parser.add_argument(
        "--deps", nargs="*", default=[],
        help="Dependency asset IDs")

    # deps
    deps_parser = subparsers.add_parser(
        "deps", help="Show/set asset dependencies")
    deps_parser.add_argument("asset_id",
                              help="Asset ID (e.g., sandcar_interior)")
    deps_parser.add_argument("dep_ids", nargs="*",
                              help="Dependency asset IDs to set "
                                   "(omit to show current)")
    deps_parser.add_argument("--clear", action="store_true",
                              help="Remove all dependencies")

    # search-corpus
    search_corpus_parser = subparsers.add_parser(
        "search-corpus",
        help="Search project corpus for reference images")
    search_corpus_parser.add_argument("episode",
                                       help="Episode number (e.g., 001)")
    search_corpus_parser.add_argument("query", nargs="?", default=None,
                                       help="Search query (default: first line of sketch script)")
    search_corpus_parser.add_argument("--json", dest="json_output",
                                       action="store_true",
                                       help="Machine-readable JSON output")

    # prepare-episode
    prep_parser = subparsers.add_parser(
        "prepare-episode",
        help="Prepare episode for orchestrator")
    prep_parser.add_argument("episode",
                             help="Episode number (e.g., 001)")
    prep_parser.add_argument("--cli", default="claude",
                             help="CLI command to use (claude, claude-zai, claude-minimax)")

    # review-episode
    rev_parser = subparsers.add_parser(
        "review-episode",
        help="Review completed episode")
    rev_parser.add_argument("episode",
                            help="Episode number (e.g., 001)")

    # new-sketch
    sketch_parser = subparsers.add_parser(
        "new-sketch",
        help="Create new sketch episode (sketch channels only)")
    sketch_parser.add_argument("premise",
                                help="Sketch premise / title")

    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        return 1

    story_path = args.story_path.resolve()

    # Override LAOZHANG_KEY from story_pipeline.json if set (per-project token)
    _sp_file = story_path / "story_pipeline.json"
    if _sp_file.exists():
        try:
            _sp = json.loads(_sp_file.read_text())
            _lk = _sp.get("laozhang_key")
            if _lk:
                os.environ["LAOZHANG_KEY"] = _lk
        except (json.JSONDecodeError, OSError):
            pass

    # Handle chat args (target + prompt parsing)
    if args.command == "chat":
        if args.prompt_extra:
            # chat S5-assets tom_morrison "prompt"
            args.target = args.prompt_or_target
            args.prompt = args.prompt_extra
        else:
            # chat S3-style "prompt"
            args.target = None
            args.prompt = args.prompt_or_target

    # Dispatch
    commands = {
        "init": cmd_init,
        "init-sketch-channel": cmd_init_sketch_channel,
        "run": cmd_run,
        "chat": cmd_chat,
        "refine": cmd_refine,
        "regenerate": cmd_regenerate,
        "status": cmd_status,
        "select": cmd_select,
        "deps": cmd_deps,
        "add-asset": cmd_add_asset,
        "search-corpus": cmd_search_corpus,
        "prepare-episode": cmd_prepare_episode,
        "review-episode": cmd_review_episode,
        "new-sketch": cmd_new_sketch,
    }

    handler = commands.get(args.command)
    if not handler:
        print(f"Unknown command: {args.command}")
        return 1

    return handler(story_path, args)


if __name__ == "__main__":
    sys.exit(main())
