from __future__ import annotations

from django.utils import timezone

from agents.steward import StewardService
from graph.models import GraphApproval, GraphApprovalStatus
from graph.native_runtime import GraphExecutionContext
from graph.registry import NodeHandlerRegistry, NodeResult
from graph.spec import ExecutionGraphSpec, GraphEdgeSpec, GraphNodeSpec


def steward_run_graph_v1() -> ExecutionGraphSpec:
    spec = ExecutionGraphSpec(
        name="steward_run",
        version=1,
        graph_type="STEWARD_RUN",
        entry="prepare",
        nodes={
            "prepare": GraphNodeSpec("prepare", "steward_prepare"),
            "collect_signals": GraphNodeSpec("collect_signals", "steward_collect_signals"),
            "run_checks": GraphNodeSpec("run_checks", "steward_run_checks"),
            "normalize_findings": GraphNodeSpec("normalize_findings", "steward_normalize_findings"),
            "classify": GraphNodeSpec("classify", "steward_classify"),
            "route": GraphNodeSpec("route", "steward_route"),
            "await_approval": GraphNodeSpec("await_approval", "steward_await_approval"),
            "summarize": GraphNodeSpec("summarize", "steward_summarize"),
            "complete": GraphNodeSpec("complete", "complete", {"terminal": True}),
        },
        edges=[
            GraphEdgeSpec("prepare", "collect_signals", "success"),
            GraphEdgeSpec("collect_signals", "run_checks", "success"),
            GraphEdgeSpec("run_checks", "normalize_findings", "success"),
            GraphEdgeSpec("normalize_findings", "classify", "success"),
            GraphEdgeSpec("classify", "route", "success"),
            GraphEdgeSpec("route", "await_approval", "approval_required"),
            GraphEdgeSpec("route", "summarize", "success"),
            GraphEdgeSpec("await_approval", "summarize", "approved"),
            GraphEdgeSpec("await_approval", "summarize", "rejected"),
            GraphEdgeSpec("summarize", "complete", "success"),
        ],
        terminal_nodes=["complete"],
        metadata={"description": "Steward V1 governance workflow; detection, classification, routing, no code mutation."},
    )
    spec.validate()
    return spec


class StewardNode:
    idempotent = True
    replay_safe = True
    destructive = False

    def __init__(self, service: StewardService, node_type: str) -> None:
        self.service = service
        self.node_type = node_type

    def steward_run_id(self, context: GraphExecutionContext) -> str:
        return str(context.graph_run.metadata["steward_run_id"])


class StewardPrepareNode(StewardNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        return NodeResult("COMPLETE", "success", {"steward_run_id": self.steward_run_id(context)})


class StewardCollectSignalsNode(StewardNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        return NodeResult("COMPLETE", "success")


class StewardRunChecksNode(StewardNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        from control_plane.projects.models import StewardRun

        steward_run = StewardRun.objects.get(id=self.steward_run_id(context))
        checks = self.service.run_checks(steward_run)
        return NodeResult("COMPLETE", "success", {"check_count": len(checks)})


class StewardNormalizeFindingsNode(StewardNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        from control_plane.projects.models import StewardRun

        steward_run = StewardRun.objects.get(id=self.steward_run_id(context))
        findings = self.service.normalize_findings(steward_run)
        return NodeResult("COMPLETE", "success", {"finding_count": len(findings)})


class StewardClassifyNode(StewardNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        from control_plane.projects.models import StewardRun

        steward_run = StewardRun.objects.get(id=self.steward_run_id(context))
        findings = self.service.classify_findings(steward_run)
        return NodeResult("COMPLETE", "success", {"classified_count": len(findings)})


class StewardRouteNode(StewardNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        from control_plane.projects.models import StewardRun

        steward_run = StewardRun.objects.get(id=self.steward_run_id(context))
        actions = self.service.route_findings(steward_run)
        approval_required = any(action.requires_approval for action in actions)
        return NodeResult("COMPLETE", "approval_required" if approval_required else "success", {"action_count": len(actions), "approval_required": approval_required})


class StewardApprovalNode(StewardNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        node_run = context.graph_run.node_runs.filter(node_id=context.graph_run.current_node).order_by("-visit_index").first()
        if GraphApproval.objects.filter(graph_run=context.graph_run, status=GraphApprovalStatus.APPROVED).exists():
            from control_plane.projects.models import StewardRun

            steward_run = StewardRun.objects.get(id=self.steward_run_id(context))
            for action in steward_run.findings.filter(actions__requires_approval=True).values_list("actions__id", flat=True).distinct():
                from control_plane.projects.models import StewardAction

                steward_action = StewardAction.objects.get(id=action)
                if steward_action.task_id is None and steward_action.action_type == "REPAIR":
                    steward_action.task = self.service._create_repair_task(steward_action.finding)
                steward_action.status = "ROUTED"
                steward_action.approved_at = timezone.now()
                steward_action.save(update_fields=["task", "status", "approved_at", "updated_at"])
            return NodeResult("COMPLETE", "approved")
        if GraphApproval.objects.filter(graph_run=context.graph_run, status=GraphApprovalStatus.REJECTED).exists():
            return NodeResult("COMPLETE", "rejected")
        GraphApproval.objects.get_or_create(graph_run=context.graph_run, node_run=node_run, reason="AWAITING_STEWARD_ROUTING_APPROVAL")
        return NodeResult("PAUSED", "awaiting", pause_reason="AWAITING_STEWARD_ROUTING_APPROVAL")


class StewardSummarizeNode(StewardNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        from control_plane.projects.models import StewardRun

        steward_run = StewardRun.objects.get(id=self.steward_run_id(context))
        self.service.complete_run(steward_run)
        return NodeResult("COMPLETE", "success", {"summary": steward_run.summary})


def steward_registry(service: StewardService) -> NodeHandlerRegistry:
    registry = NodeHandlerRegistry()
    registry.register(StewardPrepareNode(service, "steward_prepare"))
    registry.register(StewardCollectSignalsNode(service, "steward_collect_signals"))
    registry.register(StewardRunChecksNode(service, "steward_run_checks"))
    registry.register(StewardNormalizeFindingsNode(service, "steward_normalize_findings"))
    registry.register(StewardClassifyNode(service, "steward_classify"))
    registry.register(StewardRouteNode(service, "steward_route"))
    registry.register(StewardApprovalNode(service, "steward_await_approval"))
    registry.register(StewardSummarizeNode(service, "steward_summarize"))
    return registry