#!/usr/bin/env python3
"""
backfill_storyboard.py — Write spatial data and triptych prompts back to Recoil storyboards.

This is Phase 7 — intentionally last because it benefits from having real generated
frames and confirmed prompt patterns to evaluate.

Backfills:
1. Spatial data for EP002-060 (only EP001 has it) — infer camera_side, screen_direction,
   blocking from shot sequence analysis
2. triptych_prompt fields (currently all null) — pre-compute from prompt_engine
3. Wardrobe reference linking — promote best generated frames to wardrobe refs

Usage:
    python -m tools.backfill_storyboard --episode 2 --dry-run
    python -m tools.backfill_storyboard --episode 2 --write
    python -m tools.backfill_storyboard --all --dry-run
"""

import argparse
import json
import logging
import sys
from pathlib import Path
from typing import Optional

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

from recoil.pipeline._lib.recoil_bridge import (
    load_storyboard, load_breakdown, load_project_config,
    resolve_character_for_episode, get_all_scenes, RECOIL_ROOT,
)
from recoil.pipeline._lib.prompt_engine import build_cinematic_prompt, build_two_character_prompt

# ── Constants ────────────────────────────────────────────────────────

logger = logging.getLogger("starsend.backfill")

# Default spatial values when inferring from shot context
_DEFAULT_CAMERA_SIDE = "A"

# Shot types that typically have specific spatial patterns
_ANGLE_VARIANTS = {
    "CU": {"camera_side": "A", "blocking": {}},
    "ECU": {"camera_side": "A", "blocking": {}},
    "MCU": {"camera_side": "A", "blocking": {}},
    "MS": {"camera_side": "A", "blocking": {}},
    "WIDE": {"camera_side": "A", "blocking": {}},
    "LS": {"camera_side": "A", "blocking": {}},
    "INSERT": {"camera_side": "A", "blocking": {}},
}


def infer_spatial_data(shot: dict, prev_shot: Optional[dict] = None,
                        scene_shots: list[dict] = None) -> dict:
    """Infer spatial data from shot context when explicit data is missing.

    Uses the shot type, characters, and sequence position to infer
    camera_side, screen_direction, and blocking.

    This is deterministic but less precise than the storyboard agent's
    spatial analysis. Good enough for pipeline routing; will be refined
    when the storyboard agent runs on these episodes.
    """
    spatial = shot.get("spatial", {})

    # If spatial data already exists, return it as-is
    if spatial and spatial.get("camera_side") and spatial.get("blocking"):
        return spatial

    shot_type = shot.get("shot_type", "MS").upper()
    chars = shot.get("characters_in_shot", [])

    # Start with defaults from angle variant
    result = dict(_ANGLE_VARIANTS.get(shot_type, _ANGLE_VARIANTS["MS"]))

    # Infer camera side from sequence context
    if prev_shot:
        prev_spatial = prev_shot.get("spatial", {})
        prev_side = prev_spatial.get("camera_side", "A")
        # Alternate camera sides for shot/reverse-shot patterns
        if prev_shot.get("characters_in_shot") != chars and len(chars) > 0:
            result["camera_side"] = "B" if prev_side == "A" else "A"
        else:
            result["camera_side"] = prev_side

    # Infer blocking from characters
    if len(chars) == 1:
        result["blocking"] = {
            chars[0]: {"position": "center", "facing": "toward-camera"}
        }
    elif len(chars) == 2:
        result["blocking"] = {
            chars[0]: {"position": "left", "facing": "right"},
            chars[1]: {"position": "right", "facing": "left"},
        }
    elif len(chars) > 2:
        # Distribute characters across 3 reliable positions (LEFT/CENTER/RIGHT only)
        positions = ["left", "center", "right"]
        for i, char in enumerate(chars):
            pos = positions[min(i, len(positions) - 1)]
            result["blocking"][char] = {"position": pos, "facing": "toward-camera"}

    # Infer screen direction from camera movement
    movement = shot.get("camera_movement", "")
    if movement:
        if any(w in movement.lower() for w in ("left", "pan left", "track left")):
            result["screen_direction"] = "right-to-left"
        elif any(w in movement.lower() for w in ("right", "pan right", "track right")):
            result["screen_direction"] = "left-to-right"
        else:
            result["screen_direction"] = "left-to-right"
    else:
        result["screen_direction"] = "left-to-right"

    return result


def build_triptych_prompt(shot: dict, storyboard: dict,
                           episode: int, project: str = None) -> str:
    """Pre-compute the triptych/production prompt for a shot.

    Uses build_cinematic_prompt from the prompt engine, which already handles
    all the Gemini-validated patterns (kinetic descriptors, wide-shot branching,
    ENV sanitization, lighting vectors, etc).
    """
    if project is None:
        from recoil.core.paths import DEFAULT_PROJECT
        project = DEFAULT_PROJECT
    chars = shot.get("characters_in_shot", [])
    is_env = len(chars) == 0
    config = load_project_config(project)

    # Resolve character data
    character_data = {}
    for char_key in chars:
        try:
            resolved = resolve_character_for_episode(char_key, episode, project)
            character_data[char_key.lower()] = {
                "name": resolved["display_name"],
                "visual": resolved["visual_description"],
                "wardrobe": resolved["wardrobe_desc"],
                "identity_type": "non_human" if "android" in resolved.get(
                    "visual_description", "").lower() else "human",
            }
        except (FileNotFoundError, KeyError):
            pass

    if len(chars) >= 2:
        char_a = character_data.get(chars[0].lower(), {})
        char_b = character_data.get(chars[1].lower(), {})
        if char_a and char_b:
            return build_two_character_prompt(shot, storyboard, char_a, char_b, config)

    return build_cinematic_prompt(shot, storyboard, character_data, config, is_env=is_env)


def backfill_episode(episode: int, project: str = None,
                      dry_run: bool = True, write: bool = False) -> dict:
    """Backfill spatial data and triptych prompts for an episode.

    Args:
        episode: Episode number.
        project: Recoil project name.
        dry_run: If True, show changes without writing.
        write: If True, write changes to the storyboard file.

    Returns:
        Summary dict with change counts.
    """
    if project is None:
        from recoil.core.paths import DEFAULT_PROJECT
        project = DEFAULT_PROJECT
    try:
        sb = load_storyboard(episode, project)
    except FileNotFoundError:
        logger.warning(f"EP{episode:03d} storyboard not found, skipping")
        return {"episode": episode, "status": "not_found", "changes": 0}

    shots = sb.get("shots", [])
    changes = 0
    spatial_fills = 0
    prompt_fills = 0

    prev_shot = None
    for shot in shots:
        shot_id = shot.get("id", 0)
        shot_name = shot.get("name", f"shot_{shot_id}")
        modified = False

        # 1. Backfill spatial data
        existing_spatial = shot.get("spatial", {})
        if not existing_spatial or not existing_spatial.get("blocking"):
            scene_shots = shots  # Simplified — full scene context
            new_spatial = infer_spatial_data(shot, prev_shot, scene_shots)
            if new_spatial != existing_spatial:
                if dry_run:
                    print(f"  Shot {shot_id} ({shot_name}): spatial → {json.dumps(new_spatial)[:80]}...")
                else:
                    shot["spatial"] = new_spatial
                spatial_fills += 1
                modified = True

        # 2. Backfill triptych_prompt
        existing_prompt = shot.get("triptych_prompt")
        if not existing_prompt:
            try:
                new_prompt = build_triptych_prompt(shot, sb, episode, project)
                if new_prompt:
                    if dry_run:
                        print(f"  Shot {shot_id} ({shot_name}): triptych_prompt → {len(new_prompt)} chars")
                    else:
                        shot["triptych_prompt"] = new_prompt
                    prompt_fills += 1
                    modified = True
            except Exception as e:
                logger.warning(f"  Shot {shot_id}: prompt generation failed: {e}")

        if modified:
            changes += 1

        prev_shot = shot

    # Write back to file if requested
    if write and not dry_run and changes > 0:
        sb_path = RECOIL_ROOT / project / "storyboards" / f"storyboard_ep_{episode:03d}.json"
        sb_path.write_text(json.dumps(sb, indent=2, ensure_ascii=False), encoding="utf-8")
        logger.info(f"EP{episode:03d}: Wrote {changes} changes to {sb_path.name}")

    return {
        "episode": episode,
        "status": "ok",
        "total_shots": len(shots),
        "changes": changes,
        "spatial_fills": spatial_fills,
        "prompt_fills": prompt_fills,
    }


def main():
    parser = argparse.ArgumentParser(description="Backfill storyboard spatial data and prompts")
    parser.add_argument("--episode", type=int, help="Episode number to backfill")
    parser.add_argument("--all", action="store_true", help="Backfill all episodes")
    parser.add_argument("--dry-run", action="store_true", default=True,
                        help="Show changes without writing (default)")
    parser.add_argument("--write", action="store_true",
                        help="Actually write changes to storyboard files")
    parser.add_argument("--project", default=None)
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s",
                        datefmt="%H:%M:%S")

    if not args.episode and not args.all:
        parser.error("Specify --episode N or --all")

    if args.write:
        args.dry_run = False

    if args.project is None:
        from recoil.core.paths import DEFAULT_PROJECT
        args.project = DEFAULT_PROJECT

    episodes = []
    if args.all:
        # Scan for all storyboard files
        sb_dir = RECOIL_ROOT / args.project / "storyboards"
        if sb_dir.exists():
            for f in sorted(sb_dir.glob("storyboard_ep_*.json")):
                ep = int(f.stem.split("_")[-1])
                episodes.append(ep)
        if not episodes:
            print("No storyboard files found")
            return
    else:
        episodes = [args.episode]

    print(f"=== Storyboard Backfill ===")
    print(f"Episodes: {len(episodes)}")
    print(f"Mode: {'DRY RUN' if args.dry_run else 'WRITE'}")
    if not args.dry_run:
        print(f"WARNING: This will modify Recoil storyboard files!")
    print()

    total_changes = 0
    total_spatial = 0
    total_prompts = 0

    for ep in episodes:
        print(f"EP{ep:03d}:")
        result = backfill_episode(ep, args.project, dry_run=args.dry_run, write=args.write)

        if result["status"] == "not_found":
            print(f"  (storyboard not found)")
            continue

        total_changes += result["changes"]
        total_spatial += result["spatial_fills"]
        total_prompts += result["prompt_fills"]

        print(f"  {result['total_shots']} shots, {result['changes']} changes "
              f"({result['spatial_fills']} spatial, {result['prompt_fills']} prompts)")
        print()

    print(f"=== Summary ===")
    print(f"Total changes: {total_changes}")
    print(f"Spatial fills: {total_spatial}")
    print(f"Prompt fills: {total_prompts}")
    if args.dry_run:
        print(f"\nRe-run with --write to apply changes.")


if __name__ == "__main__":
    main()