"""Pipeline state machine - step definitions, dependency graph, and ordering."""

from typing import Dict, List, Optional, Tuple

# Step definitions per mode: (step_name, dependencies, display_name)
MODE_STEPS: Dict[str, List[Tuple[str, List[str], str]]] = {
    "audio_lyrics": [  # Mode A
        ("audio_preparing",            [],                         "音频准备"),
        ("demucs_separating",          ["audio_preparing"],        "人声分离"),
        ("lyric_generating",           ["demucs_separating"],      "歌词生成"),
        ("lyric_evaluating",           ["lyric_generating"],       "歌词评估"),
        ("music_generating",           ["lyric_evaluating"],       "音乐生成"),
        ("music_polling",              ["music_generating"],       "音乐轮询"),
        ("lyric_calibrating",          ["music_polling", "demucs_separating"], "歌词校准"),
        ("subtitle_rendering",         ["lyric_calibrating"],      "字幕渲染"),
        ("subtitle_exporting",         ["subtitle_rendering"],     "字幕导出"),
        ("character_designing",        ["lyric_calibrating"],      "角色设计"),
        ("character_image_generating", ["character_designing"],    "角色图生成"),
        ("storyboard_generating",      ["character_designing", "music_polling"], "分镜剧本"),
        ("scene_video_generating",     ["storyboard_generating", "character_image_generating"], "分镜视频生成"),
        ("scene_video_evaluating",     ["scene_video_generating"], "分镜视频评估"),
        ("scene_video_concatenating",  ["scene_video_evaluating"], "分镜视频拼接"),
        ("ktv_synthesizing",           ["scene_video_concatenating", "subtitle_rendering", "music_polling"], "KTV合成"),
    ],
    "video_lyrics": [  # Mode B
        ("video_preparing",            [],                         "视频准备"),
        ("demucs_separating",          ["video_preparing"],        "人声分离"),
        ("lyric_generating",           ["demucs_separating"],      "歌词生成"),
        ("lyric_evaluating",           ["lyric_generating"],       "歌词评估"),
        ("music_generating",           ["lyric_evaluating"],       "音乐生成"),
        ("music_polling",              ["music_generating"],       "音乐轮询"),
        ("lyric_calibrating",          ["music_polling", "demucs_separating"], "歌词校准"),
        ("subtitle_rendering",         ["lyric_calibrating"],      "字幕渲染"),
        ("subtitle_exporting",         ["subtitle_rendering"],     "字幕导出"),
        ("character_designing",        ["lyric_calibrating"],      "角色设计"),
        ("character_image_generating", ["character_designing"],    "角色图生成"),
        ("storyboard_generating",      ["character_designing", "music_polling"], "分镜剧本"),
        ("scene_video_generating",     ["storyboard_generating", "character_image_generating"], "分镜视频生成"),
        ("scene_video_evaluating",     ["scene_video_generating"], "分镜视频评估"),
        ("scene_video_concatenating",  ["scene_video_evaluating"], "分镜视频拼接"),
        ("ktv_synthesizing",           ["scene_video_concatenating", "subtitle_rendering", "music_polling"], "KTV合成"),
    ],
    "lyrics_only": [  # Mode C
        ("lyric_generating",           [],                         "歌词生成"),
        ("lyric_evaluating",           ["lyric_generating"],       "歌词评估"),
        ("music_generating",           ["lyric_evaluating"],       "音乐生成"),
        ("music_polling",              ["music_generating"],       "音乐轮询"),
        ("lyric_calibrating",          ["music_polling"],          "歌词校准"),
        ("subtitle_rendering",         ["lyric_calibrating"],      "字幕渲染"),
        ("subtitle_exporting",         ["subtitle_rendering"],     "字幕导出"),
        ("character_designing",        ["lyric_calibrating"],      "角色设计"),
        ("character_image_generating", ["character_designing"],    "角色图生成"),
        ("storyboard_generating",      ["character_designing", "music_polling"], "分镜剧本"),
        ("scene_video_generating",     ["storyboard_generating", "character_image_generating"], "分镜视频生成"),
        ("scene_video_evaluating",     ["scene_video_generating"], "分镜视频评估"),
        ("scene_video_concatenating",  ["scene_video_evaluating"], "分镜视频拼接"),
        ("ktv_synthesizing",           ["scene_video_concatenating", "subtitle_rendering", "music_polling"], "KTV合成"),
    ],
}

# Step states
STATE_PENDING = "pending"
STATE_RUNNING = "running"
STATE_COMPLETED = "completed"
STATE_FAILED = "failed"
STATE_SKIPPED = "skipped"

# Pipeline states
PIPELINE_SUBMITTED = "submitted"
PIPELINE_RUNNING = "running"
PIPELINE_COMPLETED = "completed"
PIPELINE_FAILED = "failed"
PIPELINE_PAUSED = "paused"  # waiting for user modification


def get_step_graph(mode: str) -> List[Tuple[str, List[str], str]]:
    """Get step definitions for a mode. Returns [(name, deps, display_name), ...]"""
    if mode not in MODE_STEPS:
        raise ValueError(f"Unknown mode: {mode}. Available: {list(MODE_STEPS.keys())}")
    return MODE_STEPS[mode]


def build_dependency_map(mode: str) -> Dict[str, dict]:
    """Build a dependency map for a mode.
    Returns: {step_name: {"deps": [...], "dependents": [...], "display_name": "...", "order": int}}
    """
    steps = get_step_graph(mode)
    dep_map = {}
    for i, (name, deps, display) in enumerate(steps):
        dep_map[name] = {
            "deps": list(deps),
            "dependents": [],
            "display_name": display,
            "order": i + 1,
        }
    # Build reverse mapping
    for name, info in dep_map.items():
        for dep in info["deps"]:
            if dep in dep_map:
                dep_map[dep]["dependents"].append(name)
    return dep_map


def get_cascade_rerun_steps(mode: str, from_step: str) -> List[str]:
    """Get all steps that need to be rerun when a step is modified.
    BFS from the modified step through dependents.
    Returns ordered list of step names.
    """
    dep_map = build_dependency_map(mode)
    if from_step not in dep_map:
        return []

    visited = set()
    queue = [from_step]
    result = []

    while queue:
        current = queue.pop(0)
        if current in visited:
            continue
        visited.add(current)
        result.append(current)
        for dep in dep_map.get(current, {}).get("dependents", []):
            if dep not in visited:
                queue.append(dep)

    # Sort by order
    result.sort(key=lambda s: dep_map.get(s, {}).get("order", 999))
    return result


def get_rerun_from_next(mode: str, from_step: str) -> List[str]:
    """When output is modified, rerun from the NEXT steps (dependents only, not the step itself)."""
    dep_map = build_dependency_map(mode)
    if from_step not in dep_map:
        return []

    direct_dependents = dep_map[from_step]["dependents"]
    all_steps = set()
    for d in direct_dependents:
        all_steps.update(get_cascade_rerun_steps(mode, d))

    result = list(all_steps)
    result.sort(key=lambda s: dep_map.get(s, {}).get("order", 999))
    return result