"""LangGraph compilation for deterministic workflow definitions."""

from __future__ import annotations

from typing import Annotated, Callable, Literal, TypedDict

from langgraph.graph import END, START, StateGraph
from langgraph.graph.state import CompiledStateGraph

from gyxx_flow.core.context import RunContext

from .model import StepDefinition, WorkflowDefinition

GraphStepStatus = Literal["success", "failed", "skipped"]


class GraphAttemptResult(TypedDict):
    """JSON-serializable attempt result stored in graph state."""

    attempt: int
    status: GraphStepStatus
    exit_code: int | None
    error: str | None
    reason: str | None


class GraphStepResult(TypedDict):
    """JSON-serializable step result stored in graph state."""

    step_id: str
    status: GraphStepStatus
    attempts: list[GraphAttemptResult]
    final_exit_code: int | None
    critical: bool
    reason: str | None


class WorkflowGraphState(TypedDict):
    """Serializable state passed between one workflow's LangGraph nodes."""

    workflow_id: str
    run_id: str
    business_date: str
    started_at: str
    shadow: bool
    dry_run: bool
    step_results: Annotated[
        dict[str, GraphStepResult],
        _merge_step_results,
    ]


GraphStepRunner = Callable[
    [StepDefinition, WorkflowGraphState],
    GraphStepResult,
]


def initial_workflow_graph_state(
    context: RunContext,
    *,
    dry_run: bool,
) -> WorkflowGraphState:
    """Create a JSON-serializable initial state from the public run context."""

    return {
        "workflow_id": context.workflow_id,
        "run_id": context.run_id,
        "business_date": context.business_date.isoformat(),
        "started_at": context.started_at.isoformat(),
        "shadow": context.shadow,
        "dry_run": dry_run,
        "step_results": {},
    }


def compile_workflow_graph(
    workflow: WorkflowDefinition,
    *,
    run_step: GraphStepRunner,
) -> CompiledStateGraph:
    """Compile dependencies into a real LangGraph DAG."""

    builder = StateGraph(WorkflowGraphState)
    ordered_steps = workflow.ordered_steps()
    dependent_ids = {
        dependency
        for step in ordered_steps
        for dependency in step.depends_on
    }
    for step in ordered_steps:
        builder.add_node(step.step_id, _step_node(step, run_step))
        if step.depends_on:
            builder.add_edge(list(step.depends_on), step.step_id)
        else:
            builder.add_edge(START, step.step_id)
        if step.step_id not in dependent_ids:
            builder.add_edge(step.step_id, END)
    return builder.compile()


def _step_node(
    step: StepDefinition,
    run_step: GraphStepRunner,
) -> Callable[[WorkflowGraphState], dict[str, object]]:
    def execute(state: WorkflowGraphState) -> dict[str, object]:
        result = run_step(step, state)
        return {"step_results": {step.step_id: result}}

    return execute


def _merge_step_results(
    current: dict[str, GraphStepResult],
    update: dict[str, GraphStepResult],
) -> dict[str, GraphStepResult]:
    """Merge branch-local results emitted in the same LangGraph superstep."""

    return {**current, **update}