from __future__ import annotations

from graph.spec import ExecutionGraphSpec, GraphEdgeSpec, GraphNodeSpec


TASK_EXECUTION_GRAPH_NAME = "task_execution"
TASK_EXECUTION_GRAPH_VERSION = 1


def task_execution_graph_v1() -> ExecutionGraphSpec:
    nodes = {
        "claim_task": GraphNodeSpec("claim_task", "claim_task", {"idempotent": False, "replay_safe": False}),
        "prepare_worktree": GraphNodeSpec("prepare_worktree", "prepare_worktree", {"idempotent": True, "replay_safe": True}),
        "build_context": GraphNodeSpec("build_context", "build_context", {"idempotent": True, "replay_safe": True}),
        "coder": GraphNodeSpec("coder", "coder", {"idempotent": False, "replay_safe": "checkpointed", "contains_internal_tool_loop": True}),
        "run_tests": GraphNodeSpec("run_tests", "run_tests", {"idempotent": True, "replay_safe": "checkpointed"}),
        "review": GraphNodeSpec("review", "review", {"idempotent": True, "replay_safe": "checkpointed"}),
        "judge": GraphNodeSpec("judge", "judge", {"idempotent": True, "replay_safe": "checkpointed"}),
        "commit": GraphNodeSpec("commit", "commit", {"idempotent": False, "replay_safe": "guarded", "destructive": True}),
        "retry_or_fail": GraphNodeSpec("retry_or_fail", "retry_or_fail", {"idempotent": False, "replay_safe": "checkpointed"}),
        "cleanup": GraphNodeSpec("cleanup", "cleanup", {"idempotent": True, "replay_safe": True}),
        "complete": GraphNodeSpec("complete", "complete", {"terminal": True}),
        "fail": GraphNodeSpec("fail", "fail", {"terminal": True}),
    }
    edges = [
        GraphEdgeSpec("claim_task", "prepare_worktree", "success"),
        GraphEdgeSpec("prepare_worktree", "build_context", "success"),
        GraphEdgeSpec("build_context", "coder", "success"),
        GraphEdgeSpec("coder", "run_tests", "success"),
        GraphEdgeSpec("coder", "retry_or_fail", "failure"),
        GraphEdgeSpec("run_tests", "review", "complete"),
        GraphEdgeSpec("review", "judge", "PASS"),
        GraphEdgeSpec("review", "retry_or_fail", "REWORK_REQUIRED"),
        GraphEdgeSpec("review", "retry_or_fail", "REJECTED"),
        GraphEdgeSpec("judge", "commit", "PASS"),
        GraphEdgeSpec("judge", "retry_or_fail", "FAIL"),
        GraphEdgeSpec("commit", "cleanup", "success"),
        GraphEdgeSpec("cleanup", "complete", "success"),
        GraphEdgeSpec("retry_or_fail", "build_context", "retry_available"),
        GraphEdgeSpec("retry_or_fail", "cleanup", "retry_exhausted"),
        GraphEdgeSpec("cleanup", "fail", "failed"),
    ]
    spec = ExecutionGraphSpec(
        name=TASK_EXECUTION_GRAPH_NAME,
        version=TASK_EXECUTION_GRAPH_VERSION,
        graph_type="TASK_EXECUTION",
        entry="claim_task",
        nodes=nodes,
        edges=edges,
        terminal_nodes=["complete", "fail"],
        metadata={"description": "Task execution graph v1 mirrors AutonomousTaskLoop semantics."},
    )
    spec.validate()
    return spec


def task_execution_graph_v2_static_analysis() -> ExecutionGraphSpec:
    spec = task_execution_graph_v1()
    nodes = dict(spec.nodes)
    nodes["static_analysis"] = GraphNodeSpec("static_analysis", "static_analysis", {"deterministic": True})
    edges = []
    for edge in spec.edges:
        if edge.source == "run_tests" and edge.target == "review":
            edges.append(GraphEdgeSpec("run_tests", "static_analysis", edge.condition, edge.metadata))
        elif edge.source == "review" and edge.target == "judge":
            edges.append(edge)
        else:
            edges.append(edge)
    edges.append(GraphEdgeSpec("static_analysis", "review", "PASS"))
    challenger = ExecutionGraphSpec(
        name=TASK_EXECUTION_GRAPH_NAME,
        version=2,
        graph_type=spec.graph_type,
        entry=spec.entry,
        nodes=nodes,
        edges=edges,
        terminal_nodes=spec.terminal_nodes,
        metadata={"parent_version": 1, "change_summary": "Insert deterministic static_analysis node before Reviewer."},
    )
    challenger.validate()
    return challenger