"""
Frame Index — cross-episode reference matching for NanoBanana.

Builds and queries an index of selected frames from completed episodes
so that later episodes can use them as reference images when standard
references/ files don't match asset IDs.

Usage:
    from lib.frame_index import load_frame_index, build_frame_index

    idx = build_frame_index(story_path, "001")
    fmt = format_frame_index_for_prompt(idx)
"""

import json
import logging
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)

INDEX_FILENAME = "frame_index.json"
VERSION = 1


def load_frame_index(story_path: Path) -> dict:
    """Load frame_index.json from story root.

    Returns empty structure if the file is missing or corrupt.
    """
    index_path = story_path / INDEX_FILENAME
    if not index_path.exists():
        return {"version": VERSION, "frames": []}

    try:
        data = json.loads(index_path.read_text())
        if not isinstance(data.get("frames"), list):
            logger.warning("frame_index.json has invalid 'frames' field")
            return {"version": VERSION, "frames": []}
        return data
    except (json.JSONDecodeError, OSError) as e:
        logger.warning(f"Failed to load frame_index.json: {e}")
        return {"version": VERSION, "frames": []}


def build_frame_index(story_path: Path, episode_num: str) -> dict:
    """Build/update frame index from a completed episode.

    Reads pipeline.json -> 08-frames selections, 09-script-adapted scene
    descriptions, and 04-assets appearance data.  Merges results into the
    existing frame_index.json (replaces entries for the same episode).

    No AI calls — purely file-based.
    """
    episode_path = story_path / episode_num

    # Load pipeline.json
    pipeline_file = episode_path / "pipeline.json"
    if not pipeline_file.exists():
        logger.warning(f"No pipeline.json in {episode_path}")
        return load_frame_index(story_path)

    try:
        pipeline = json.loads(pipeline_file.read_text())
    except (json.JSONDecodeError, OSError) as e:
        logger.error(f"Failed to read pipeline.json: {e}")
        return load_frame_index(story_path)

    # --- 08-frames selections ---
    frames_map = pipeline.get("08-frames", {})
    if not frames_map:
        logger.warning(f"No 08-frames in pipeline.json for episode {episode_num}")
        return load_frame_index(story_path)

    # --- 09-script-adapted scene data ---
    script_iter = pipeline.get("09-script-adapted")
    scenes_by_num: dict[int, dict] = {}
    if script_iter:
        script_file = episode_path / "09-script-adapted" / script_iter / "script.json"
        if script_file.exists():
            try:
                script_data = json.loads(script_file.read_text())
                for scene in script_data.get("scenes", []):
                    num = scene.get("number")
                    if num is not None:
                        scenes_by_num[int(num)] = scene
            except (json.JSONDecodeError, OSError) as e:
                logger.warning(f"Failed to read adapted script: {e}")

    # --- 04-assets appearance data ---
    assets_iter = pipeline.get("04-assets")
    assets_chars: dict[str, list[int]] = {}  # char_id -> appears_in
    assets_locs: dict[str, list[int]] = {}
    assets_objs: dict[str, list[int]] = {}
    if assets_iter:
        assets_file = episode_path / "04-assets" / assets_iter / "assets.json"
        if assets_file.exists():
            try:
                assets_data = json.loads(assets_file.read_text())
                for c in assets_data.get("characters", []):
                    cid = c.get("id", "")
                    if cid:
                        assets_chars[cid] = [int(s) for s in c.get("appears_in", [])]
                for loc in assets_data.get("locations", []):
                    lid = loc.get("id", "")
                    if lid:
                        assets_locs[lid] = [int(s) for s in loc.get("appears_in", [])]
                for obj in assets_data.get("objects", []):
                    oid = obj.get("id", "")
                    if oid:
                        assets_objs[oid] = [int(s) for s in obj.get("appears_in", [])]
            except (json.JSONDecodeError, OSError) as e:
                logger.warning(f"Failed to read assets.json: {e}")

    # --- Build frame entries ---
    new_frames = []
    for scene_key, frame_file in sorted(frames_map.items()):
        # Parse scene number from key like "01-scene" -> 1
        try:
            scene_num = int(scene_key.split("-")[0])
        except (ValueError, IndexError):
            logger.warning(f"Cannot parse scene number from key: {scene_key}")
            continue

        # Verify the frame file exists
        frame_path_rel = f"{episode_num}/08-frames/{scene_key}/{frame_file}.jpg"
        frame_path_abs = story_path / frame_path_rel
        if not frame_path_abs.exists():
            # Try .png
            frame_path_rel = f"{episode_num}/08-frames/{scene_key}/{frame_file}.png"
            frame_path_abs = story_path / frame_path_rel
            if not frame_path_abs.exists():
                logger.warning(f"Frame file not found: {frame_path_rel}")
                continue

        # Get description from adapted script
        scene_data = scenes_by_num.get(scene_num, {})
        description = scene_data.get("what_is_shown", "")
        shot_type = scene_data.get("shot_type", "")

        # Find which characters/locations/objects appear in this scene
        chars_in_scene = [
            cid for cid, scenes in assets_chars.items()
            if scene_num in scenes
        ]
        locs_in_scene = [
            lid for lid, scenes in assets_locs.items()
            if scene_num in scenes
        ]
        objs_in_scene = [
            oid for oid, scenes in assets_objs.items()
            if scene_num in scenes
        ]

        entry = {
            "episode": episode_num,
            "scene": scene_num,
            "scene_key": scene_key,
            "frame_file": frame_file,
            "path": frame_path_rel,
            "description": description,
            "characters": chars_in_scene,
            "locations": locs_in_scene,
            "objects": objs_in_scene,
            "shot_type": shot_type,
        }
        new_frames.append(entry)

    logger.info(f"Built {len(new_frames)} frame entries for episode {episode_num}")

    # Merge into existing index (replace same-episode entries)
    existing = load_frame_index(story_path)
    other_frames = [f for f in existing.get("frames", []) if f.get("episode") != episode_num]
    all_frames = other_frames + new_frames

    # Sort by episode, then scene
    all_frames.sort(key=lambda f: (f.get("episode", ""), f.get("scene", 0)))

    index = {
        "version": VERSION,
        "last_updated": datetime.now(timezone.utc).isoformat(),
        "frames": all_frames,
    }

    # Save
    index_path = story_path / INDEX_FILENAME
    index_path.write_text(json.dumps(index, indent=2, ensure_ascii=False) + "\n")
    logger.info(f"Saved {INDEX_FILENAME}: {len(all_frames)} total frames")

    return index


def format_frame_index_for_prompt(frame_index: dict) -> str:
    """Format indexed frames as readable text for Claude prompt injection.

    Each entry shows episode, scene, relative path, and full description
    so Claude can pick which frames to use as references.
    """
    frames = frame_index.get("frames", [])
    if not frames:
        return ""

    lines = [
        "\n\n## Available Episode Frames\n",
        "These are selected frames from completed episodes. When an asset in "
        "this episode matches something visible in a previous episode frame, "
        "add a `frame_references` field to that asset listing the frame paths "
        "you consider the best visual match. Read each description carefully "
        "and select frames where the asset is clearly visible and prominent. "
        "These frames will be sent to NanoBanana as reference images.\n",
    ]

    current_episode = None
    for f in frames:
        ep = f.get("episode", "???")
        if ep != current_episode:
            current_episode = ep
            lines.append(f"\n### Episode {ep}\n")

        scene = f.get("scene", "?")
        path = f.get("path", "")
        desc = f.get("description", "(no description)")
        shot = f.get("shot_type", "")
        chars = f.get("characters", [])
        locs = f.get("locations", [])
        objs = f.get("objects", [])

        # Truncate very long descriptions
        if len(desc) > 400:
            desc = desc[:397] + "..."

        tags = []
        if chars:
            tags.append(f"characters: {', '.join(chars)}")
        if locs:
            tags.append(f"locations: {', '.join(locs)}")
        if objs:
            tags.append(f"objects: {', '.join(objs)}")

        tag_str = f" [{'; '.join(tags)}]" if tags else ""
        shot_str = f" ({shot})" if shot else ""

        lines.append(f"- **Scene {scene}**{shot_str}: `{path}`{tag_str}")
        lines.append(f"  {desc}\n")

    lines.append(
        "\nTo reference a frame, add to the asset entry in assets.json:\n"
        '```json\n"frame_references": ["001/08-frames/16-scene/003-refined-003-refined-003.jpg"]\n```\n'
    )

    return "\n".join(lines)


def get_frames_for_asset(frame_index: dict, asset_id: str) -> list[dict]:
    """Find all frames containing a specific asset.

    Searches characters, locations, and objects lists in each frame entry.
    Returns entries sorted by relevance — exact asset_id matches in the
    primary list come first.
    """
    frames = frame_index.get("frames", [])
    if not frames:
        return []

    matches = []
    for f in frames:
        found_in = []
        if asset_id in f.get("characters", []):
            found_in.append("characters")
        if asset_id in f.get("locations", []):
            found_in.append("locations")
        if asset_id in f.get("objects", []):
            found_in.append("objects")

        if found_in:
            matches.append({
                **f,
                "_found_in": found_in,
            })

    # Sort: primary asset category first, then by episode/scene
    def sort_key(m):
        # Prefer entries where asset appears as object (most specific visual)
        priority = 2
        if "objects" in m["_found_in"]:
            priority = 0
        elif "characters" in m["_found_in"]:
            priority = 1
        return (priority, m.get("episode", ""), m.get("scene", 0))

    matches.sort(key=sort_key)
    return matches
