#!/usr/bin/env python3
"""
Stage S8: Post-Episode Review

Claude reviews completed episode assembly and updates continuity.json
for cross-episode consistency.
"""

import subprocess
import sys

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

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

from lib import save_metadata, ensure_dir
from lib.path_utils_v3 import get_next_iteration
from lib.story_pipeline_state import get_project_type


def execute(story_path: Path, episode_num: str, pipeline: dict) -> dict:
    """
    Execute Stage S8: Post-Episode Review.

    Args:
        story_path: Path to story project directory
        episode_num: Episode number string (e.g., "001")
        pipeline: Story pipeline state dict
    """
    is_sketch = get_project_type(pipeline) == "sketch_channel"
    label = "Sketch" if is_sketch else "Episode"
    print(f"\nStage S8: Review {label} {episode_num}")

    episode_path = story_path / episode_num
    if not episode_path.exists():
        return {"status": "error", "error": f"Episode directory not found: {episode_path}"}

    # Check if assembly exists — try both 13-assembly and 12-assembly
    assembly_dir = None
    for asm_name in ["13-assembly", "12-assembly"]:
        candidate = episode_path / asm_name
        if candidate.exists() and list(candidate.glob("*/final.mp4")):
            assembly_dir = candidate
            break
    if not assembly_dir:
        return {"status": "error",
                "error": f"{label} {episode_num} has no assembled video. "
                         f"Complete the episode pipeline first."}

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

    # Sketch channels may use channel_bible.json
    if is_sketch:
        bible_rel = f"S2-bible/{bible_iter}/channel_bible.json"
        if not (story_path / bible_rel).exists():
            bible_rel = f"S2-bible/{bible_iter}/story_bible.json"
    else:
        bible_rel = f"S2-bible/{bible_iter}/story_bible.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()

    # Load current continuity
    continuity_file = story_path / "continuity.json"
    if continuity_file.exists():
        continuity_state = continuity_file.read_text()
    else:
        continuity_state = f"(No continuity state yet — this is the first {label.lower()} review)"

    # Build prompt — choose sketch or story template
    prompts_dir = Path(__file__).parent.parent.parent / "prompts" / "story"
    if is_sketch:
        prompt_path = prompts_dir / "stage_s8_review_sketch.txt"
    else:
        prompt_path = prompts_dir / "stage_s8_review.txt"

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

    ep_number = int(episode_num)

    # Get episode title
    episode_title = ""
    if is_sketch:
        # For sketch channels, get title from brief.txt or sketch_script.md
        for candidate in [episode_path / "brief.txt",
                          episode_path / "sketch_script.md",
                          episode_path / "01-story" / "001" / "story.md"]:
            if candidate.exists():
                first_line = candidate.read_text().strip().split("\n")[0]
                episode_title = first_line.lstrip("# ").strip()
                if episode_title.upper().startswith(("SKETCH", "EPISODE")):
                    parts = episode_title.split(":", 1)
                    episode_title = parts[1].strip() if len(parts) > 1 else episode_title
                break
    else:
        episodes_iter = pipeline.get("S6-episodes")
        if episodes_iter:
            episodes_file = story_path / "S6-episodes" / episodes_iter / "episodes.json"
            if episodes_file.exists():
                eps_data = json.loads(episodes_file.read_text())
                for ep in eps_data.get("episodes", []):
                    if ep.get("number") == ep_number:
                        episode_title = ep.get("title", "")
                        break

    prompt = prompt_path.read_text()
    prompt = prompt.replace("{episode_number}", str(ep_number))
    prompt = prompt.replace("{episode_title}", episode_title)
    prompt = prompt.replace("{episode_path}", episode_num)
    prompt = prompt.replace("{bible_path}", bible_rel)
    prompt = prompt.replace("{continuity_state}", continuity_state)

    # Output — write to episode directory
    prompt = (f"## Output Directory\n\n"
              f"Write continuity_update.json and user_message.txt to: {episode_num}/\n\n"
              f"{prompt}")

    # Call Claude
    print(f"Session: {session_id}")
    print(f"Reviewing {label.lower()} {ep_number}: \"{episode_title}\"")

    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}: {error_detail}"}
    except subprocess.TimeoutExpired:
        return {"status": "error", "error": "Claude timed out"}

    # Process continuity update
    update_file = episode_path / "continuity_update.json"
    if update_file.exists():
        update = json.loads(update_file.read_text())
        if is_sketch:
            _merge_sketch_continuity(story_path, update)
        else:
            _merge_continuity(story_path, update)
        print("Updated: continuity.json")
    else:
        print("Warning: No continuity_update.json created")

    # Build frame index for cross-episode reference matching
    from lib.frame_index import build_frame_index
    try:
        frame_idx = build_frame_index(story_path, episode_num)
        n_frames = len(frame_idx.get("frames", []))
        print(f"Updated: frame_index.json ({n_frames} frames)")
    except Exception as e:
        print(f"Warning: Failed to build frame index: {e}")

    # User message
    user_msg_file = episode_path / "user_message.txt"
    review_msg_file = episode_path / "review_user_message.txt"
    if review_msg_file.exists():
        user_message = review_msg_file.read_text().strip()
    elif user_msg_file.exists():
        user_message = user_msg_file.read_text().strip()
    else:
        user_message = f"{label} {ep_number} reviewed. Continuity updated."

    print(f"\nReviewed: {label} {episode_num}")

    return {
        "status": "success",
        "episode": episode_num,
        "continuity_updated": update_file.exists(),
        "user_message": user_message,
    }


def _merge_sketch_continuity(story_path: Path, update: dict) -> None:
    """Merge sketch continuity update — tracks behavior patterns + running gags."""
    continuity_file = story_path / "continuity.json"

    if continuity_file.exists():
        try:
            existing = json.loads(continuity_file.read_text())
        except (json.JSONDecodeError, OSError):
            existing = {}
    else:
        existing = {}

    # Last sketch reviewed
    existing["last_sketch_reviewed"] = update.get(
        "last_sketch_reviewed",
        existing.get("last_sketch_reviewed"))

    # Merge characters (including behavior_patterns, catchphrases)
    if "characters" not in existing:
        existing["characters"] = {}
    for char_id, char_data in update.get("characters", {}).items():
        if char_id not in existing["characters"]:
            existing["characters"][char_id] = {}
        ex_char = existing["characters"][char_id]
        # Simple fields: overwrite
        for key in ["visual_notes", "consistency_tips", "last_seen_in"]:
            if key in char_data:
                ex_char[key] = char_data[key]
        # List fields: merge unique values
        for key in ["behavior_patterns", "catchphrases"]:
            if key in char_data:
                existing_list = ex_char.get(key, [])
                for item in char_data[key]:
                    if item not in existing_list:
                        existing_list.append(item)
                ex_char[key] = existing_list

    # Merge locations
    if "locations" not in existing:
        existing["locations"] = {}
    for loc_id, loc_data in update.get("locations", {}).items():
        if loc_id not in existing["locations"]:
            existing["locations"][loc_id] = {}
        existing["locations"][loc_id].update(loc_data)

    # Merge running gags
    if "running_gags" not in existing:
        existing["running_gags"] = {}
    for gag_id, gag_data in update.get("running_gags", {}).items():
        if gag_id in existing["running_gags"]:
            # Update existing gag
            ex_gag = existing["running_gags"][gag_id]
            ex_gag["times_used"] = gag_data.get("times_used",
                                                  ex_gag.get("times_used", 0))
            if "usage_notes" in gag_data:
                ex_gag["usage_notes"] = gag_data["usage_notes"]
            if "potential" in gag_data:
                ex_gag["potential"] = gag_data["potential"]
        else:
            existing["running_gags"][gag_id] = gag_data

    # Replace quality notes (latest wins)
    if "quality_notes" in update:
        existing["quality_notes"] = update["quality_notes"]

    # Append new references
    if "new_references" in update:
        if "new_references" not in existing:
            existing["new_references"] = []
        existing["new_references"].extend(update["new_references"])

    # Save
    with open(continuity_file, "w") as f:
        json.dump(existing, f, indent=2)
        f.write("\n")


def _merge_continuity(story_path: Path, update: dict) -> None:
    """Merge continuity update into the main continuity.json."""
    continuity_file = story_path / "continuity.json"

    if continuity_file.exists():
        try:
            existing = json.loads(continuity_file.read_text())
        except (json.JSONDecodeError, OSError):
            existing = {}
    else:
        existing = {}

    # Deep merge
    existing["last_episode_reviewed"] = update.get("last_episode_reviewed",
                                                     existing.get("last_episode_reviewed"))

    # Merge characters
    if "characters" not in existing:
        existing["characters"] = {}
    for char_id, char_data in update.get("characters", {}).items():
        if char_id not in existing["characters"]:
            existing["characters"][char_id] = {}
        existing["characters"][char_id].update(char_data)

    # Merge locations
    if "locations" not in existing:
        existing["locations"] = {}
    for loc_id, loc_data in update.get("locations", {}).items():
        if loc_id not in existing["locations"]:
            existing["locations"][loc_id] = {}
        existing["locations"][loc_id].update(loc_data)

    # Replace story progress (latest wins)
    if "story_progress" in update:
        existing["story_progress"] = update["story_progress"]

    # Replace quality notes (latest wins)
    if "quality_notes" in update:
        existing["quality_notes"] = update["quality_notes"]

    # Append new references
    if "new_references" in update:
        if "new_references" not in existing:
            existing["new_references"] = []
        existing["new_references"].extend(update["new_references"])

    # Save
    with open(continuity_file, "w") as f:
        json.dump(existing, f, indent=2)
        f.write("\n")
