#!/usr/bin/env python3
"""
check_consistency.py — Run cross-shot character consistency checks.

Usage:
    python3 tools/check_consistency.py --project afterimage [--character Sadie]
    python3 tools/check_consistency.py --project afterimage --all

Scans all generated frames, groups by character (from shot plans),
and runs the grid-then-pairs consistency check for each character.
"""

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

# Ensure imports work
sys.path.insert(0, str(Path(__file__).parent.parent))
sys.path.insert(0, str(Path(__file__).parent.parent / "pipeline"))

from recoil.core.paths import projects_root, ProjectPaths
from visual.elements import get_identity_anchor

logger = logging.getLogger(__name__)


def find_character_frames(project: str, character_name: str = "") -> dict[str, dict[str, Path]]:
    """Find all generated frames grouped by character.

    Returns: {character_name: {shot_id: frame_path}}
    """
    project_dir = projects_root() / project
    plans_dir = ProjectPaths.for_project(project).plans_dir
    frames_dir = project_dir / "output" / "frames"

    if not plans_dir.exists():
        logger.error("No plans directory: %s", plans_dir)
        return {}

    # Read all shot plans to find character-to-shot mapping
    character_frames: dict[str, dict[str, Path]] = {}

    for plan_file in sorted(plans_dir.glob("ep_*_plan.json")):
        try:
            plan = json.loads(plan_file.read_text())
        except (json.JSONDecodeError, OSError) as e:
            logger.warning("Could not read plan %s: %s", plan_file, e)
            continue

        shots = plan.get("shots", [])
        ep_num = plan_file.stem.replace("_plan", "")

        for shot in shots:
            shot_id = shot.get("shot_id", shot.get("id", ""))
            characters_in_shot = shot.get("characters", [])

            # Find the frame file
            ep_frames = frames_dir / ep_num
            if not ep_frames.exists():
                continue

            # Look for the frame (various naming conventions)
            # Use delimiter-bounded patterns to avoid substring matches
            # (e.g., SH1 must not match SH10, SH11, etc.)
            frame_path = None
            for pattern in [f"*{shot_id}.*", f"*{shot_id}_*", f"shot_{shot_id}.*", f"shot_{shot_id}_*", f"sh{shot_id}.*", f"sh{shot_id}_*"]:
                matches = list(ep_frames.glob(pattern))
                if matches:
                    # Prefer .png, then .jpg
                    frame_path = next(
                        (m for m in matches if m.suffix == ".png"),
                        matches[0]
                    )
                    break

            if frame_path is None:
                continue

            # Map character -> shot -> frame
            for char in characters_in_shot:
                char_name = char if isinstance(char, str) else char.get("name", char.get("char_id", ""))
                if character_name and char_name.lower() != character_name.lower():
                    continue
                if char_name not in character_frames:
                    character_frames[char_name] = {}
                character_frames[char_name][f"{ep_num}_{shot_id}"] = frame_path

    return character_frames


def main():
    parser = argparse.ArgumentParser(description="Cross-shot character consistency check")
    parser.add_argument("--project", required=True, help="Project name")
    parser.add_argument("--character", default="", help="Check specific character (default: all)")
    parser.add_argument("--all", action="store_true", help="Check all characters")
    parser.add_argument("-v", "--verbose", action="store_true")
    args = parser.parse_args()

    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(levelname)s: %(message)s"
    )

    # Load global bible for identity anchors
    bible_path = ProjectPaths.for_project(args.project).global_bible_path
    global_bible = {}
    if bible_path.exists():
        try:
            global_bible = json.loads(bible_path.read_text())
        except json.JSONDecodeError:
            logger.warning("Could not parse global bible")

    # Find frames grouped by character
    character_frames = find_character_frames(args.project, args.character)

    if not character_frames:
        print("No character frames found.")
        return

    # Import the consistency checker
    from recoil.pipeline._lib.critics.character_consistency_critic import check_character_consistency

    # Run checks
    all_consistent = True
    for char_name, shot_frames in sorted(character_frames.items()):
        anchor = get_identity_anchor(global_bible, char_name)
        print(f"\n{'='*60}")
        print(f"Character: {char_name} ({len(shot_frames)} shots)")
        if anchor:
            print(f"Anchor: {anchor[:80]}...")
        print(f"{'='*60}")

        result = check_character_consistency(
            character_name=char_name,
            shot_frames=shot_frames,
            character_anchor=anchor,
        )

        if result["consistent"]:
            print(f"  CONSISTENT across all {result['total_shots']} shots")
        else:
            all_consistent = False
            print(f"  INCONSISTENT — {len(result['outlier_shots'])} outlier(s):")
            for detail in result["details"]:
                if not detail.get("passed", True):
                    print(f"    Shot {detail['shot_id']} (vs reference {detail['reference_shot']}):")
                    for fc in detail.get("failed_checks", []):
                        print(f"      - {fc['name']}: {fc['answer']} [{fc['severity']}]")

    print(f"\n{'='*60}")
    if all_consistent:
        print("ALL CHARACTERS CONSISTENT")
    else:
        print("INCONSISTENCIES FOUND — review flagged shots")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()