#!/usr/bin/env python3
"""
Stage S7: Prepare Episode

Creates episode directory with rich brief.txt and generates episode title card.
Runs per-episode, BEFORE the orchestrator pipeline.
"""

import subprocess
import sys

from lib import claude_env
import json
import shutil
from pathlib import Path
from datetime import datetime, timezone

sys.path.insert(0, str(Path(__file__).parent.parent))

repo_root = Path(__file__).parent.parent.parent.parent

from lib import save_metadata, ensure_dir
from lib.path_utils_v3 import get_next_iteration
from lib.pipeline import PIPELINE_STAGES, load_pipeline
from lib.story_pipeline_state import get_project_type


def execute(story_path: Path, episode_num: str, pipeline: dict,
            cli_command: str = "claude") -> dict:
    """
    Execute Stage S7: Prepare Episode.

    Args:
        story_path: Path to story project directory
        episode_num: Episode number string (e.g., "001")
        pipeline: Story pipeline state dict
        cli_command: CLI to use ("claude", "claude-zai", "claude-minimax")
    """
    if get_project_type(pipeline) == "sketch_channel":
        return _execute_sketch(story_path, episode_num, pipeline, cli_command)
    return _execute_story(story_path, episode_num, pipeline, cli_command)


def _execute_story(story_path: Path, episode_num: str, pipeline: dict,
                   cli_command: str = "claude") -> dict:
    """Prepare episode for story adaptation mode."""
    print(f"\nStage S7: Prepare Episode {episode_num}")

    # Resolve story-level artifacts
    bible_iter = pipeline.get("S2-bible")
    style_iter = pipeline.get("S3-style")
    assets_done = pipeline.get("S5-assets")
    episodes_iter = pipeline.get("S6-episodes")

    if not all([bible_iter, style_iter, episodes_iter]):
        return {"status": "error",
                "error": "Story stages S2, S3, S6 must be complete before preparing episodes"}

    bible_path = story_path / "S2-bible" / bible_iter / "story_bible.json"
    style_path = story_path / "S3-style" / style_iter / "style.json"
    episodes_path = story_path / "S6-episodes" / episodes_iter / "episodes.json"

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

    # Load episode data
    episodes_data = json.loads(episodes_path.read_text())
    ep_number = int(episode_num)
    episode = None
    for ep in episodes_data.get("episodes", []):
        if ep.get("number") == ep_number:
            episode = ep
            break

    if not episode:
        return {"status": "error",
                "error": f"Episode {ep_number} not found in episodes.json"}

    # 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()

    # Create episode directory and orchestrator structure
    episode_path = story_path / episode_num
    ensure_dir(episode_path)

    # Create pipeline.json for episode (orchestrator format)
    load_pipeline(episode_path)

    # Create stage directories for orchestrator
    for stage in PIPELINE_STAGES:
        ensure_dir(episode_path / stage)
    ensure_dir(episode_path / "01-story" / "001")

    # Build continuity section
    continuity_section = _build_continuity_section(story_path, ep_number)

    # Build character references section
    char_refs_section = _build_character_references(story_path, episode)

    # Load bible for additional context
    bible = json.loads(bible_path.read_text())
    style = json.loads(style_path.read_text())

    # Determine mode: session-resume (claude) or standalone (claude-zai, etc.)
    use_standalone = cli_command != "claude"

    # Select prompt template
    prompts_dir = Path(__file__).parent.parent.parent / "prompts" / "story"
    if use_standalone:
        prompt_path = prompts_dir / "stage_s7_prepare_standalone.txt"
    else:
        prompt_path = prompts_dir / "stage_s7_prepare.txt"

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

    prompt = prompt_path.read_text()

    # Common substitutions
    prompt = prompt.replace("{bible_path}",
                            f"S2-bible/{bible_iter}/story_bible.json")
    prompt = prompt.replace("{style_path}",
                            f"S3-style/{style_iter}/style.json")
    prompt = prompt.replace("{episodes_path}",
                            f"S6-episodes/{episodes_iter}/episodes.json")
    prompt = prompt.replace("{episode_number}", str(ep_number))
    prompt = prompt.replace("{episode_title}", episode.get("title", ""))
    prompt = prompt.replace("{continuity_section}", continuity_section)
    prompt = prompt.replace("{character_references_section}", char_refs_section)

    # Standalone-specific substitutions
    if use_standalone:
        source_lines = episode.get("source_text_range", "")
        # Convert "Lines 20–112" -> "20-112" for template
        import re
        lines_match = re.search(r"(\d+)\s*[–-]\s*(\d+)", source_lines)
        source_lines_clean = f"{lines_match.group(1)}-{lines_match.group(2)}" if lines_match else source_lines
        prompt = prompt.replace("{source_lines}", source_lines_clean)
        prompt = prompt.replace("{output_dir}", episode_num)
        prompt = prompt.replace("{episode_title_upper}", episode.get("title", "").upper())

    # Inject corpus references if available
    corpus_section = _build_corpus_refs_section(episode_path)
    if corpus_section:
        prompt += "\n" + corpus_section
        print(f"  Added corpus references to prompt")

    # Output directory prefix (session mode only — standalone has it in template)
    if not use_standalone:
        prompt = f"## Output Directory\n\nWrite brief.txt to: {episode_num}/\nWrite other files to: {episode_num}/\n\n{prompt}"

    return _run_prepare_claude(story_path, episode_path, episode_num,
                                ep_number, episode, pipeline,
                                cli_command, session_id, prompt,
                                use_standalone)


def _execute_sketch(story_path: Path, episode_num: str, pipeline: dict,
                    cli_command: str = "claude") -> dict:
    """Prepare episode for sketch channel mode."""
    print(f"\nStage S7: Prepare Sketch {episode_num}")

    # Resolve story-level artifacts — sketch mode doesn't need S6-episodes
    bible_iter = pipeline.get("S2-bible")
    style_iter = pipeline.get("S3-style")

    if not all([bible_iter, style_iter]):
        return {"status": "error",
                "error": "Channel stages S2 (bible), S3 (style) must be complete before preparing sketches"}

    bible_path = story_path / "S2-bible" / bible_iter / "channel_bible.json"
    # Fallback to story_bible.json if channel_bible.json doesn't exist
    if not bible_path.exists():
        bible_path = story_path / "S2-bible" / bible_iter / "story_bible.json"
    style_path = story_path / "S3-style" / style_iter / "style.json"

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

    ep_number = int(episode_num)

    # Sketch script: check episode dir for sketch_script.md or 01-story/001/story.md
    episode_path = story_path / episode_num
    sketch_script_path = None
    for candidate in [
        episode_path / "sketch_script.md",
        episode_path / "01-story" / "001" / "story.md",
    ]:
        if candidate.exists():
            sketch_script_path = str(candidate.relative_to(story_path))
            break

    if not sketch_script_path:
        return {"status": "error",
                "error": f"No sketch script found in {episode_num}/. "
                         f"Create {episode_num}/sketch_script.md or use new-sketch command."}

    # Get sketch title from script first line
    script_text = (story_path / sketch_script_path).read_text().strip()
    first_line = script_text.split("\n")[0]
    episode_title = first_line.lstrip("# ").strip()

    # 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()

    # Create episode directory and orchestrator structure
    ensure_dir(episode_path)
    load_pipeline(episode_path)
    for stage in PIPELINE_STAGES:
        ensure_dir(episode_path / stage)
    ensure_dir(episode_path / "01-story" / "001")

    # Build continuity section (sketch mode: behavior patterns + running gags)
    continuity_section = _build_sketch_continuity_section(story_path, ep_number)

    # Build character references (use all characters from bible)
    char_refs_section = _build_character_references_from_bible(
        story_path, bible_path)

    # Select sketch prompt template (always standalone-style for sketches)
    prompts_dir = Path(__file__).parent.parent.parent / "prompts" / "story"
    prompt_path = prompts_dir / "stage_s7_prepare_sketch.txt"
    if not prompt_path.exists():
        return {"status": "error", "error": f"Sketch prompt template not found: {prompt_path}"}

    prompt = prompt_path.read_text()

    # Substitutions
    prompt = prompt.replace("{bible_path}",
                            str(bible_path.relative_to(story_path)))
    prompt = prompt.replace("{style_path}",
                            f"S3-style/{style_iter}/style.json")
    prompt = prompt.replace("{sketch_script_path}", sketch_script_path)
    prompt = prompt.replace("{episode_number}", str(ep_number))
    prompt = prompt.replace("{episode_title}", episode_title)
    prompt = prompt.replace("{episode_title_upper}", episode_title.upper())
    prompt = prompt.replace("{continuity_section}", continuity_section)
    prompt = prompt.replace("{character_references_section}", char_refs_section)
    prompt = prompt.replace("{output_dir}", episode_num)

    creative_notes = pipeline.get("creative_notes")
    if creative_notes:
        prompt += (f"\n\n## CREATIVE DIRECTION (from channel config)\n\n"
                   f"Follow these as hard requirements:\n\n{creative_notes}\n")

    # Inject corpus references if available
    corpus_section = _build_corpus_refs_section(episode_path)
    if corpus_section:
        prompt += "\n" + corpus_section
        print(f"  Added corpus references to prompt")

    use_standalone = cli_command != "claude"

    # Output directory prefix (session mode)
    if not use_standalone:
        prompt = (f"## Output Directory\n\n"
                  f"Write brief.txt to: {episode_num}/\n"
                  f"Write other files to: {episode_num}/\n\n{prompt}")

    episode = {"title": episode_title}
    return _run_prepare_claude(story_path, episode_path, episode_num,
                                ep_number, episode, pipeline,
                                cli_command, session_id, prompt,
                                use_standalone)


def _run_prepare_claude(story_path: Path, episode_path: Path,
                        episode_num: str, ep_number: int,
                        episode: dict, pipeline: dict,
                        cli_command: str, session_id: str,
                        prompt: str, use_standalone: bool) -> dict:
    """Run Claude to generate brief.txt and finalize episode directory."""
    # Build CLI command
    print(f"CLI: {cli_command}" + (f" (session: {session_id})" if not use_standalone else " (standalone)"))
    print(f"Preparing episode {ep_number}: \"{episode.get('title', '')}\"")

    if use_standalone:
        cmd = [cli_command,
               "--allowedTools", "Read,Write,Glob,Grep",
               "-p", prompt]
    else:
        cmd = [cli_command, "--resume", session_id,
               "--allowedTools", "Read,Write,Glob,Grep",
               "-p", prompt]

    try:
        result = subprocess.run(
            cmd,
            capture_output=True, text=True,
            cwd=str(story_path), timeout=1200,
            env=claude_env()
        )
        if result.returncode != 0:
            error_detail = result.stderr or result.stdout
            return {"status": "error",
                    "error": f"{cli_command} returned code {result.returncode}: {error_detail}"}
    except subprocess.TimeoutExpired:
        return {"status": "error", "error": f"{cli_command} timed out"}

    # Validate brief.txt was created
    brief_file = episode_path / "brief.txt"
    if not brief_file.exists():
        return {"status": "error", "error": "brief.txt not created for episode"}

    # Copy brief.txt as story.md into 01-story/001/ — this IS the story,
    # so stage 01 auto-selects it and skips redundant generation
    story_md = episode_path / "01-story" / "001" / "story.md"
    ensure_dir(story_md.parent)
    shutil.copy(brief_file, story_md)
    print(f"  brief.txt -> 01-story/001/story.md (stage 01 will auto-skip)")

    # User message
    user_msg_file = episode_path / "user_message.txt"
    user_message = (user_msg_file.read_text().strip() if user_msg_file.exists()
                    else f"Episode {ep_number} prepared. Run orchestrator to start production.")

    print(f"\nPrepared: {episode_num}/")
    print(f"  brief.txt ({brief_file.stat().st_size} bytes)")
    print(f"  01-story/001/story.md (auto-skip stage 01)")
    print(f"\nNext: orchestrator_v5.py {story_path} --episode {episode_num} run")

    return {
        "status": "success",
        "episode": episode_num,
        "title": episode.get("title", ""),
        "user_message": user_message,
    }


def _build_corpus_refs_section(episode_path: Path) -> str:
    """Build corpus references section for prepare prompt, if corpus_refs.json exists."""
    refs_file = episode_path / "corpus_refs.json"
    if not refs_file.exists():
        return ""

    try:
        data = json.loads(refs_file.read_text())
    except (json.JSONDecodeError, OSError):
        return ""

    results = data.get("results", [])
    if not results:
        return ""

    lines = [
        "\n## Reference Corpus Results\n",
        "These reference images were found by searching the project corpus.",
        "Each is available at the listed path. Use them as composition/style guides.\n",
    ]

    for r in results:
        local_ref = r.get("local_ref")
        if not local_ref:
            continue
        desc = r.get("description", "")
        if len(desc) > 200:
            desc = desc[:200] + "..."
        lines.append(f"{r['rank']}. **{r['title']}** (score: {r['score']:.3f})")
        if desc:
            lines.append(f"   Description: {desc}")
        lines.append(f"   Image: `{local_ref}`")
        lines.append("")

    lines.extend([
        "## Scene-to-Reference Assignment\n",
        "For each scene, optionally assign ONE corpus reference as the compositional foundation.",
        "Write the assignment in each scene's \"Corpus ref:\" field.",
        "After writing brief.txt, write updated corpus_refs.json with scene_assignments filled.",
        "Not every scene needs a corpus ref — only assign where a strong match exists.",
    ])

    return "\n".join(lines)


def _build_continuity_section(story_path: Path, current_episode: int) -> str:
    """Build continuity notes from previous episodes."""
    continuity_file = story_path / "continuity.json"
    if not continuity_file.exists() or current_episode <= 1:
        return "(First episode — no continuity notes yet)"

    try:
        continuity = json.loads(continuity_file.read_text())
    except (json.JSONDecodeError, OSError):
        return "(No continuity data available)"

    sections = ["## Continuity from Previous Episodes\n"]

    # Character visual state
    chars = continuity.get("characters", {})
    if chars:
        sections.append("### Character Visual State")
        for char_id, char_data in chars.items():
            notes = char_data.get("visual_notes", "")
            tips = char_data.get("consistency_tips", "")
            if notes or tips:
                sections.append(f"**{char_id}**: {notes}")
                if tips:
                    sections.append(f"  Consistency: {tips}")

    # Location state
    locs = continuity.get("locations", {})
    if locs:
        sections.append("\n### Location Visual State")
        for loc_id, loc_data in locs.items():
            notes = loc_data.get("visual_notes", "")
            if notes:
                sections.append(f"**{loc_id}**: {notes}")

    # Story progress
    progress = continuity.get("story_progress", {})
    if progress:
        sections.append("\n### Story Progress")
        if progress.get("shown"):
            sections.append(f"Previously shown: {progress['shown']}")
        if progress.get("next_setup"):
            sections.append(f"Setup for this episode: {progress['next_setup']}")

    # Quality notes
    quality = continuity.get("quality_notes", {})
    if quality:
        worked = quality.get("worked_well", [])
        adjust = quality.get("adjust_next", [])
        if worked:
            sections.append("\n### What Worked")
            for w in worked:
                sections.append(f"- {w}")
        if adjust:
            sections.append("\n### Adjustments Needed")
            for a in adjust:
                sections.append(f"- {a}")

    return "\n".join(sections)


def _build_sketch_continuity_section(story_path: Path, current_episode: int) -> str:
    """Build continuity notes for sketch channel (behavior patterns + running gags)."""
    continuity_file = story_path / "continuity.json"
    if not continuity_file.exists() or current_episode <= 1:
        return "(First sketch — no continuity notes yet)"

    try:
        continuity = json.loads(continuity_file.read_text())
    except (json.JSONDecodeError, OSError):
        return "(No continuity data available)"

    sections = ["## Continuity from Previous Sketches\n"]

    # Character behavior patterns
    chars = continuity.get("characters", {})
    if chars:
        sections.append("### Character Notes")
        for char_id, char_data in chars.items():
            notes = char_data.get("visual_notes", "")
            tips = char_data.get("consistency_tips", "")
            patterns = char_data.get("behavior_patterns", [])
            catchphrases = char_data.get("catchphrases", [])
            parts = [f"**{char_id}**:"]
            if notes:
                parts.append(f"  Visual: {notes}")
            if tips:
                parts.append(f"  Consistency: {tips}")
            if patterns:
                parts.append(f"  Behavior: {', '.join(patterns)}")
            if catchphrases:
                parts.append(f"  Catchphrases: {', '.join(repr(c) for c in catchphrases)}")
            sections.append("\n".join(parts))

    # Running gags
    gags = continuity.get("running_gags", {})
    if gags:
        sections.append("\n### Running Gags")
        for gag_id, gag_data in gags.items():
            desc = gag_data.get("description", "")
            potential = gag_data.get("potential", "")
            times = gag_data.get("times_used", 0)
            sections.append(f"- **{gag_id}**: {desc} (used {times}x, {potential})")

    # Location state
    locs = continuity.get("locations", {})
    if locs:
        sections.append("\n### Location Notes")
        for loc_id, loc_data in locs.items():
            notes = loc_data.get("visual_notes", "")
            if notes:
                sections.append(f"**{loc_id}**: {notes}")

    # Quality notes
    quality = continuity.get("quality_notes", {})
    if quality:
        worked = quality.get("worked_well", [])
        adjust = quality.get("adjust_next", [])
        if worked:
            sections.append("\n### What Worked")
            for w in worked:
                sections.append(f"- {w}")
        if adjust:
            sections.append("\n### Adjustments Needed")
            for a in adjust:
                sections.append(f"- {a}")

    return "\n".join(sections)


def _build_character_references_from_bible(story_path: Path,
                                            bible_path: Path) -> str:
    """Build character references for all characters in the channel bible."""
    refs_path = story_path / "references" / "people"

    try:
        bible = json.loads(bible_path.read_text())
    except (json.JSONDecodeError, OSError):
        return "(Could not load channel bible)"

    # Get character IDs from bible
    chars = bible.get("characters", [])
    if not chars:
        return "(No characters in channel bible)"

    sections = []
    for char in chars:
        char_id = char.get("id", "")
        if not char_id:
            continue

        char_dir = refs_path / char_id
        parts = [f"### {char_id} — {char.get('name', '')}"]

        # Description from references/people/{id}/description.md
        desc_file = char_dir / "description.md"
        if desc_file.exists():
            parts.append(desc_file.read_text().strip())

        # Reference images
        sheet = char_dir / "sheet.jpg"
        portrait = char_dir / "portrait.jpg"
        if sheet.exists():
            parts.append(f"Reference sheet: `{sheet}`")
        elif portrait.exists():
            parts.append(f"Portrait: `{portrait}`")

        if not desc_file.exists() and not char_dir.exists():
            parts.append("(No reference generated yet)")

        sections.append("\n".join(parts))

    return "\n\n".join(sections) if sections else "(No character data)"


def _build_character_references(story_path: Path, episode: dict) -> str:
    """Build character reference info for episode characters."""
    refs_path = story_path / "references" / "people"
    if not refs_path.exists():
        return "(No character references generated yet)"

    chars = episode.get("characters", [])
    if not chars:
        return "(No characters in this episode)"

    sections = []
    for char_id in chars:
        char_dir = refs_path / char_id
        if not char_dir.exists():
            sections.append(f"**{char_id}**: No reference generated")
            continue

        desc_file = char_dir / "description.md"
        sheet = char_dir / "sheet.jpg"
        portrait = char_dir / "portrait.jpg"

        parts = [f"### {char_id}"]
        if desc_file.exists():
            parts.append(desc_file.read_text())
        if sheet.exists():
            parts.append(f"Reference sheet: `{sheet}`")
        elif portrait.exists():
            parts.append(f"Portrait: `{portrait}`")

        sections.append("\n".join(parts))

    return "\n\n".join(sections) if sections else "(No character data)"


