"""Build a canonical PlanGraph from DPE stage records."""

from __future__ import annotations

from dpe.engine.states import STAGE_ARTIFACT_MAP, STAGE_ORDER, Stage
from dpe.models.manifest import RunManifest

from ._run_id import to_core_run_id

# Dependency DAG encoding the actual DPE data flow.
_STAGE_DEPS: dict[Stage, list[Stage]] = {
    Stage.INPUT_NORMALIZE: [],
    Stage.IDEATION_A: [Stage.INPUT_NORMALIZE],
    Stage.IDEATION_B: [Stage.INPUT_NORMALIZE],
    Stage.CROSS_CRITIQUE_A: [Stage.IDEATION_A, Stage.IDEATION_B],
    Stage.CROSS_CRITIQUE_B: [Stage.IDEATION_A, Stage.IDEATION_B],
    Stage.SYNTHESIS: [
        Stage.IDEATION_A,
        Stage.IDEATION_B,
        Stage.CROSS_CRITIQUE_A,
        Stage.CROSS_CRITIQUE_B,
    ],
    Stage.PRD_V1: [Stage.SYNTHESIS],
    Stage.PRD_CRITIQUE_R1: [Stage.PRD_V1],
    Stage.PRD_REVISION_R1: [Stage.PRD_V1, Stage.PRD_CRITIQUE_R1],
    Stage.PRD_CRITIQUE_R2: [Stage.PRD_REVISION_R1],
    Stage.PRD_REVISION_R2: [Stage.PRD_REVISION_R1, Stage.PRD_CRITIQUE_R2],
    Stage.FINALIZE: [Stage.PRD_REVISION_R2],
}

# If PRD_REVISION_R2 was skipped, FINALIZE falls back to these in order.
_FINALIZE_FALLBACKS = [Stage.PRD_REVISION_R1, Stage.PRD_V1, Stage.SYNTHESIS]

_STAGE_ROLES: dict[Stage, str] = {
    Stage.INPUT_NORMALIZE: "normalizer",
    Stage.IDEATION_A: "ideator_a",
    Stage.IDEATION_B: "ideator_b",
    Stage.CROSS_CRITIQUE_A: "critic_a",
    Stage.CROSS_CRITIQUE_B: "critic_b",
    Stage.SYNTHESIS: "synthesizer",
    Stage.PRD_V1: "prd_author",
    Stage.PRD_CRITIQUE_R1: "prd_critic",
    Stage.PRD_REVISION_R1: "prd_author",
    Stage.PRD_CRITIQUE_R2: "prd_critic",
    Stage.PRD_REVISION_R2: "prd_author",
    Stage.FINALIZE: "editor",
}


def _node_id(stage: Stage) -> str:
    return f"dpe-{stage.value.lower()}"


def _resolve_finalize_deps(manifest: RunManifest) -> list[str]:
    """Determine FINALIZE dependencies based on which stages actually ran."""
    executed = {r.stage for r in manifest.stages if not r.skipped}
    if Stage.PRD_REVISION_R2.value in executed:
        return [_node_id(Stage.PRD_REVISION_R2)]
    for fallback in _FINALIZE_FALLBACKS:
        if fallback.value in executed:
            return [_node_id(fallback)]
    return []


def _stage_status(stage: Stage, manifest: RunManifest) -> str:
    """Map a stage to its plan-graph status."""
    for rec in manifest.stages:
        if rec.stage == stage.value:
            if rec.skipped:
                return "skipped"
            if rec.error:
                return "failed"
            return "completed"

    # Stage not in records at all — check if the run stopped before reaching it
    if manifest.stop_reason != "completed":
        executed_stages = {r.stage for r in manifest.stages}
        stage_idx = STAGE_ORDER.index(stage)
        # If any earlier stage was recorded, this one is pending (never reached)
        if any(
            STAGE_ORDER.index(Stage(s)) < stage_idx
            for s in executed_stages
            if s in [st.value for st in STAGE_ORDER]
        ):
            return "pending"
    return "pending"


def build_plan_graph(manifest: RunManifest) -> dict:
    """Build a canonical plan_graph dict from a DPE RunManifest.

    Returns a dict conforming to plan-graph/1.0 schema.
    """
    core_run_id = to_core_run_id(manifest.run_id)

    # Build lookup: stage_name -> StageRecord
    stage_records = {r.stage: r for r in manifest.stages}

    nodes = []
    for stage in STAGE_ORDER:
        rec = stage_records.get(stage.value)

        deps = _STAGE_DEPS[stage]
        if stage == Stage.FINALIZE:
            dep_ids = _resolve_finalize_deps(manifest)
        else:
            dep_ids = [_node_id(d) for d in deps]

        node: dict = {
            "node_id": _node_id(stage),
            "name": stage.value,
            "status": _stage_status(stage, manifest),
            "depends_on": dep_ids,
            "agent": rec.adapter_used if rec else "none",
            "metadata": {
                "dpe_artifact": STAGE_ARTIFACT_MAP.get(stage, ""),
                "dpe_role": _STAGE_ROLES.get(stage, "unknown"),
            },
        }
        nodes.append(node)

    return {
        "schema_version": "plan-graph/1.0",
        "run_id": core_run_id,
        "description": "DPE dialectical pipeline",
        "nodes": nodes,
    }
