"""
Pipeline state management for V5 orchestrator.

pipeline.json is the single source of truth for which iteration is selected
at each stage. No symlinks.
"""

import json
from pathlib import Path
from typing import Optional


# Ordered list of pipeline stages
PIPELINE_STAGES = [
    "01-story",
    "02-understanding",
    "03-script",
    "04-assets",
    "05-references",
    "06-characteristic",
    "07-grids",
    "08-frames",
    "09-script-adapted",
    "10-videos",
    "11-titles",
    "12-captions",
    "13-assembly",
]

# Which parent stages each stage reads from
STAGE_DEPENDENCIES = {
    "01-story":        [],
    "02-understanding": ["01-story"],
    "03-script":         ["02-understanding"],
    "04-assets":         ["03-script"],
    "05-references":     ["04-assets"],
    "06-characteristic": ["03-script", "04-assets"],
    "07-grids":          ["03-script", "04-assets", "06-characteristic"],
    "08-frames":         ["07-grids"],
    "09-script-adapted": ["03-script", "08-frames"],
    "10-videos":         ["08-frames", "06-characteristic", "09-script-adapted"],
    "11-titles":         ["10-videos"],
    "12-captions":       ["03-script", "10-videos"],
    "13-assembly":       ["10-videos", "12-captions", "11-titles"],
}


def load_pipeline(take_path: Path) -> dict:
    """Load pipeline.json. Create with all-null if missing.

    Includes migration logic for old stage names.
    """
    pipeline_file = take_path / "pipeline.json"

    if pipeline_file.exists():
        with open(pipeline_file) as f:
            data = json.load(f)

        # Migration: rename old stage keys
        key_renames = {
            "08-script-adapted": "09-script-adapted",
            "09-frames": "08-frames",
            "11-captions": "12-captions",
            "12-assembly": "13-assembly",
        }

        migrated = False
        for old_key, new_key in key_renames.items():
            if old_key in data and new_key not in data:
                data[new_key] = data.pop(old_key)
                migrated = True

        # Add new keys if missing
        for stage in PIPELINE_STAGES:
            if stage not in data:
                data[stage] = None
                migrated = True

        # Rename on-disk directories to match new stage names
        # Must use temp names to avoid collision (08↔09 swap)
        dir_renames = {
            "08-script-adapted": "09-script-adapted",
            "09-frames": "08-frames",
            "11-captions": "12-captions",
            "12-assembly": "13-assembly",
        }

        # Check if any old dirs exist that need renaming
        pending_renames = {}
        for old_dir, new_dir in dir_renames.items():
            old_path = take_path / old_dir
            new_path = take_path / new_dir
            if old_path.exists() and not new_path.exists():
                pending_renames[old_dir] = new_dir

        if pending_renames:
            # Use temp names to handle the 08↔09 swap
            import tempfile
            temp_map = {}
            for old_dir in pending_renames:
                old_path = take_path / old_dir
                tmp_name = f"_tmp_migrate_{old_dir}"
                tmp_path = take_path / tmp_name
                old_path.rename(tmp_path)
                temp_map[tmp_name] = pending_renames[old_dir]

            for tmp_name, new_dir in temp_map.items():
                tmp_path = take_path / tmp_name
                new_path = take_path / new_dir
                tmp_path.rename(new_path)

            migrated = True

        if migrated:
            save_pipeline(take_path, data)

        return data

    # Create default
    pipeline = {stage: None for stage in PIPELINE_STAGES}
    save_pipeline(take_path, pipeline)
    return pipeline


def save_pipeline(take_path: Path, pipeline: dict) -> None:
    """Save pipeline.json."""
    pipeline_file = take_path / "pipeline.json"
    with open(pipeline_file, "w") as f:
        json.dump(pipeline, f, indent=2)
        f.write("\n")


def find_next_stage(pipeline: dict, take_path: Path = None) -> Optional[str]:
    """
    Find first stage not completed.

    For regular stages: checks if None (not started)
    For scene-based stages (08, 10, 11): checks if all scenes selected

    Args:
        pipeline: Pipeline dict
        take_path: Optional take path (needed for scene count validation)

    Returns:
        Next stage name or None if complete
    """
    for stage in PIPELINE_STAGES:
        value = pipeline.get(stage)

        # Scene-based stages
        if stage in ["08-frames", "10-videos"]:
            if not isinstance(value, dict):
                return stage

            # Check if all scenes selected
            if take_path and (take_path / "08-frames" / "extraction_metadata.json").exists():
                import json
                with open(take_path / "08-frames" / "extraction_metadata.json") as f:
                    meta = json.load(f)
                    scene_count = meta["scene_count"]

                selected_count = sum(1 for v in value.values() if v)
                if selected_count < scene_count:
                    return stage
            else:
                if not value or not all(value.values()):
                    return stage
            continue

        # Titles stage: scene-based with 2 fixed scenes (intro, outro)
        if stage == "11-titles":
            if not isinstance(value, dict):
                return stage
            if not value.get("intro") or not value.get("outro"):
                return stage
            continue

        # Regular stages: check if None
        if value is None:
            return stage

    return None


def resolve_inputs(pipeline: dict, stage: str, take_path=None) -> dict:
    """
    Get parent stage iterations from pipeline for a given stage.

    Returns e.g. {"03-script": "006", "04-assets": "004"}
    Raises ValueError if any required parent is null.

    Special handling:
    - Stage 10: Returns per-scene frame selections from 08-frames
    - Stage 11-titles: Returns per-scene video selections
    - Stage 12/13: Returns per-scene video selections
    """
    # Special case: Stage 10 needs per-scene frame selections
    if stage == "10-videos":
        if take_path is None:
            raise ValueError("take_path required for Stage 10 input resolution")

        # Get frame selections from main pipeline
        frame_selections = pipeline.get("08-frames")

        # Check frames are a dict (selections made)
        if not isinstance(frame_selections, dict):
            raise ValueError(
                "Frames not extracted/selected. Run: extract-frames and select-frame"
            )

        # Get scene count from extraction metadata
        frames_dir = take_path / "08-frames"
        extraction_meta_file = frames_dir / "extraction_metadata.json"
        if not extraction_meta_file.exists():
            raise ValueError("Extraction metadata not found")

        import json
        with open(extraction_meta_file) as f:
            extraction_meta = json.load(f)
        scene_count = extraction_meta["scene_count"]

        # Validate ALL scenes have selections
        unselected = []
        for i in range(1, scene_count + 1):
            scene_key = f"{i:02d}-scene"
            if not frame_selections.get(scene_key):
                unselected.append(scene_key)

        if unselected:
            raise ValueError(
                f"Scenes not selected: {', '.join(unselected)}\n"
                f"Run: select-frame SCENE FRAME for each unselected scene"
            )

        # Return frame selections dict + other dependencies
        return {
            "08-frames": frame_selections,  # Dict: scene -> frame filename
            "06-characteristic": pipeline.get("06-characteristic"),
            "09-script-adapted": pipeline.get("09-script-adapted"),
        }

    # Special case: Stages that need per-scene video selections
    if stage in ["11-titles", "12-captions", "13-assembly"]:
        if take_path is None:
            raise ValueError(f"take_path required for {stage} input resolution")

        # Get video selections from main pipeline
        video_selections = pipeline.get("10-videos")

        # Check videos are a dict (selections made)
        if not isinstance(video_selections, dict):
            raise ValueError(
                "Videos not generated/selected. Run: new-video and select-video"
            )

        # Get scene count from frame extraction metadata
        frames_dir = take_path / "08-frames"
        extraction_meta_file = frames_dir / "extraction_metadata.json"
        if not extraction_meta_file.exists():
            raise ValueError("Extraction metadata not found")

        import json
        with open(extraction_meta_file) as f:
            extraction_meta = json.load(f)
        scene_count = extraction_meta["scene_count"]

        # Validate ALL scenes have video selections
        unselected = []
        for i in range(1, scene_count + 1):
            scene_key = f"{i:02d}-scene"
            if not video_selections.get(scene_key):
                unselected.append(scene_key)

        if unselected:
            raise ValueError(
                f"Scenes not selected: {', '.join(unselected)}\n"
                f"Run: select-video SCENE ITERATION for each unselected scene"
            )

        # Return video selections dict + other standard dependencies
        result = {
            "10-videos": video_selections,  # Dict: scene -> iteration
        }

        # Add other dependencies from STAGE_DEPENDENCIES
        deps = STAGE_DEPENDENCIES.get(stage, [])
        for dep in deps:
            if dep != "10-videos":  # Already handled above
                if dep == "11-titles":
                    # Title selections are a dict (intro/outro)
                    title_selections = pipeline.get("11-titles")
                    result[dep] = title_selections
                else:
                    parent_iter = pipeline.get(dep)
                    if parent_iter is None:
                        raise ValueError(
                            f"Cannot generate {stage}: parent stage {dep} "
                            f"has no selected iteration"
                        )
                    result[dep] = parent_iter

        return result

    # Special case: 09-script-adapted needs frame selections
    if stage == "09-script-adapted":
        deps = STAGE_DEPENDENCIES.get(stage, [])
        inputs = {}
        for parent_stage in deps:
            parent_val = pipeline.get(parent_stage)
            if parent_val is None:
                raise ValueError(
                    f"Cannot generate {stage}: parent stage {parent_stage} "
                    f"has no selected iteration"
                )
            inputs[parent_stage] = parent_val
        return inputs

    # Standard resolution for other stages
    deps = STAGE_DEPENDENCIES.get(stage, [])
    inputs = {}

    for parent_stage in deps:
        parent_iter = pipeline.get(parent_stage)
        if parent_iter is None:
            raise ValueError(
                f"Cannot generate {stage}: parent stage {parent_stage} "
                f"has no selected iteration"
            )
        inputs[parent_stage] = parent_iter

    return inputs


def get_downstream_stages(stage: str) -> list:
    """Get all stages that depend (directly or transitively) on the given stage."""
    downstream = []
    stage_idx = PIPELINE_STAGES.index(stage) if stage in PIPELINE_STAGES else -1

    for s in PIPELINE_STAGES[stage_idx + 1:]:
        deps = STAGE_DEPENDENCIES.get(s, [])
        # Direct dependency
        if stage in deps:
            downstream.append(s)
        # Transitive: if any of this stage's deps are already in downstream
        elif any(d in downstream for d in deps):
            downstream.append(s)

    return downstream
