"""
Cascading stage regeneration for V4 orchestrator.

When user wants to go back to an earlier stage and regenerate all dependent stages.
"""

from pathlib import Path
from .feedback import append_feedback


# Canonical stage ordering with task templates.
# Must match init_story_v4.py task queue initialization.
STAGE_PIPELINE = [
    {"stage": "02-understanding", "tasks": [{"action": "generate"}, {"action": "await_feedback"}]},
    {"stage": "03-script",        "tasks": [{"action": "generate"}, {"action": "await_feedback"}]},
    {"stage": "04-assets",        "tasks": [{"action": "generate"}, {"action": "await_feedback"}]},
    {"stage": "05-references",    "tasks": [{"action": "check_skip"}, {"action": "search_or_skip"}]},
    {"stage": "06-characteristic", "tasks": [{"action": "generate"}, {"action": "await_feedback"}]},
    {"stage": "07-grids",         "tasks": [{"action": "generate"}, {"action": "await_feedback"}]},
]

# Quick lookup: stage name -> index in pipeline
_STAGE_INDEX = {s["stage"]: i for i, s in enumerate(STAGE_PIPELINE)}


def get_stage_index(stage_name: str) -> int:
    """Get index of stage in pipeline. Raises KeyError if not found."""
    if stage_name not in _STAGE_INDEX:
        raise KeyError(f"Unknown stage: {stage_name}. Known: {list(_STAGE_INDEX.keys())}")
    return _STAGE_INDEX[stage_name]


def get_stages_between(from_stage: str, to_stage: str) -> list:
    """
    Get stage definitions between from_stage and to_stage (inclusive).

    Returns list of stage dicts from STAGE_PIPELINE.
    """
    from_idx = get_stage_index(from_stage)
    to_idx = get_stage_index(to_stage)

    if from_idx > to_idx:
        raise ValueError(f"from_stage ({from_stage}) is after to_stage ({to_stage})")

    return STAGE_PIPELINE[from_idx:to_idx + 1]


def get_available_restart_stages(current_stage: str) -> list:
    """
    Get list of stages user can restart from (all stages before current).

    Returns list of {"stage": "...", "label": "..."} dicts.
    """
    current_idx = get_stage_index(current_stage)
    stages = []
    for stage_def in STAGE_PIPELINE[:current_idx]:
        stages.append({
            "stage": stage_def["stage"],
            "label": stage_def["stage"].split("-", 1)[1].replace("-", " ").title()
        })
    return stages


def create_cascade_tasks(
    queue: list,
    feedback_task: dict,
    target_stage: str,
    current_stage: str,
    feedback_text: str,
    take_path: Path,
    logger
) -> None:
    """
    Create cascade of tasks from target_stage through current_stage.

    1. Mark feedback_task completed with feedback
    2. Log feedback for target_stage (so Claude sees it in history)
    3. Create new generate+feedback tasks for each stage in range
    4. Rewire downstream deps to point to new chain's final task

    Args:
        queue: Task queue (modified in place)
        feedback_task: The await_feedback task that triggered the restart
        target_stage: Stage to restart from (e.g., "02-understanding")
        current_stage: Stage user is currently at (e.g., "06-characteristic")
        feedback_text: User's feedback text
        take_path: Path to take directory
        logger: Generation logger
    """
    # Edge case: restart_from same stage = just regenerate
    if target_stage == current_stage:
        logger.info(f"restart_from same stage ({target_stage}), treating as regenerate")
        return  # Caller should fall through to regenerate logic

    # Validate stage ordering
    target_idx = get_stage_index(target_stage)
    current_idx = get_stage_index(current_stage)
    if target_idx >= current_idx:
        logger.error(f"Cannot restart_from {target_stage} — it's not before {current_stage}")
        return

    # 1. Mark current feedback task as completed
    feedback_task["status"] = "completed"
    feedback_task["feedback"] = f"restart_from:{target_stage}: {feedback_text}"
    logger.info(f"Task {feedback_task['id']} marked completed (restart_from {target_stage})")

    # 2. Log feedback for the target stage so Claude sees it
    append_feedback(
        take_path,
        stage=target_stage,
        iteration="cascade",
        feedback_type="restart_from",
        text=feedback_text,
        action=f"cascade_from_{current_stage}"
    )
    logger.info(f"Logged feedback for {target_stage}")

    # 3. Get stages in cascade range
    stages_in_range = get_stages_between(target_stage, current_stage)
    cascade_stage_names = {s["stage"] for s in stages_in_range}
    logger.info(f"Cascade range: {[s['stage'] for s in stages_in_range]}")

    # 4. Create new tasks
    max_id = max(t["id"] for t in queue)
    prev_task_id = feedback_task["id"]
    new_tasks = []

    for stage_def in stages_in_range:
        stage_name = stage_def["stage"]

        for task_template in stage_def["tasks"]:
            max_id += 1
            new_task = {
                "id": max_id,
                "stage": stage_name,
                "action": task_template["action"],
                "status": "pending",
                "deps": [prev_task_id],
                "iteration": None,
                "parent_iteration": None,
                "result": None,
                "message": None,
                "feedback": None,
                "error": None
            }

            # First generate task carries the feedback text
            if stage_name == target_stage and task_template["action"] == "generate":
                new_task["message"] = f"Cascade from {current_stage}: {feedback_text[:100]}"

            new_tasks.append(new_task)
            prev_task_id = max_id

    # Insert after the feedback task
    insert_index = queue.index(feedback_task) + 1
    for i, new_task in enumerate(new_tasks):
        queue.insert(insert_index + i, new_task)

    new_final_task_id = new_tasks[-1]["id"]
    new_task_ids = {t["id"] for t in new_tasks}

    logger.info(f"Created {len(new_tasks)} new tasks (IDs {new_tasks[0]['id']}-{new_tasks[-1]['id']})")

    # 5. Rewire downstream deps
    rewire_downstream_deps(queue, cascade_stage_names, new_final_task_id, new_task_ids, logger)


def rewire_downstream_deps(
    queue: list,
    cascade_stage_names: set,
    new_final_task_id: int,
    new_task_ids: set,
    logger
) -> None:
    """
    Rewire pending downstream tasks to depend on the new cascade's final task.

    Scans all pending tasks NOT in the cascade range. If their deps point to
    any old task in a cascade stage, rewire to new_final_task_id.
    """
    # Collect all OLD task IDs belonging to cascade stages
    old_cascade_task_ids = set()
    for t in queue:
        if t["stage"] in cascade_stage_names and t["id"] not in new_task_ids:
            old_cascade_task_ids.add(t["id"])

    # Rewire pending tasks in stages AFTER the cascade
    for t in queue:
        if t["status"] != "pending":
            continue
        if t["id"] in new_task_ids:
            continue
        if t["stage"] in cascade_stage_names:
            continue

        old_deps = t.get("deps", [])
        new_deps = []
        changed = False

        for dep_id in old_deps:
            if dep_id in old_cascade_task_ids:
                if new_final_task_id not in new_deps:
                    new_deps.append(new_final_task_id)
                changed = True
            else:
                new_deps.append(dep_id)

        if changed:
            t["deps"] = new_deps
            logger.info(f"Rewired task {t['id']} ({t['stage']}/{t['action']}) deps: {old_deps} -> {new_deps}")
