from __future__ import annotations

from django.utils import timezone

from agents.lifecycle import EvolutionService, ExplorerService, ExtensionService
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 project_extension_graph_v1() -> ExecutionGraphSpec:
    spec = ExecutionGraphSpec(
        name="project_extension",
        version=1,
        graph_type="PROJECT_EXTENSION",
        entry="prepare",
        nodes={node: GraphNodeSpec(node, node if node == "complete" else f"extension_{node}") for node in ["prepare", "gather_context", "plan_with_project_brain", "validate_plan", "await_approval", "materialize_project_dag", "execute", "verify_extension", "complete"]},
        edges=[GraphEdgeSpec("prepare", "gather_context", "success"), GraphEdgeSpec("gather_context", "plan_with_project_brain", "success"), GraphEdgeSpec("plan_with_project_brain", "validate_plan", "success"), GraphEdgeSpec("validate_plan", "await_approval", "success"), GraphEdgeSpec("await_approval", "materialize_project_dag", "approved"), GraphEdgeSpec("await_approval", "complete", "rejected"), GraphEdgeSpec("materialize_project_dag", "execute", "success"), GraphEdgeSpec("execute", "verify_extension", "success"), GraphEdgeSpec("verify_extension", "complete", "success")],
        terminal_nodes=["complete"],
        metadata={"description": "Project extension workflow: Sol plan, approval, normal task execution, milestone verification."},
    )
    spec.validate()
    return spec


def project_evolution_graph_v1() -> ExecutionGraphSpec:
    spec = ExecutionGraphSpec(
        name="project_evolution",
        version=1,
        graph_type="PROJECT_EVOLUTION",
        entry="prepare",
        nodes={node: GraphNodeSpec(node, node if node in ["complete", "not_improved"] else f"evolution_{node}") for node in ["prepare", "establish_baseline", "form_hypothesis", "plan_with_project_brain", "await_approval", "materialize_work", "execute", "measure_candidate", "compare_baseline", "judge_evolution", "complete", "not_improved"]},
        edges=[GraphEdgeSpec("prepare", "establish_baseline", "success"), GraphEdgeSpec("establish_baseline", "form_hypothesis", "success"), GraphEdgeSpec("form_hypothesis", "plan_with_project_brain", "success"), GraphEdgeSpec("plan_with_project_brain", "await_approval", "success"), GraphEdgeSpec("await_approval", "materialize_work", "approved"), GraphEdgeSpec("await_approval", "not_improved", "rejected"), GraphEdgeSpec("materialize_work", "execute", "success"), GraphEdgeSpec("execute", "measure_candidate", "success"), GraphEdgeSpec("measure_candidate", "compare_baseline", "success"), GraphEdgeSpec("compare_baseline", "judge_evolution", "success"), GraphEdgeSpec("judge_evolution", "complete", "PASS"), GraphEdgeSpec("judge_evolution", "not_improved", "NOT_IMPROVED")],
        terminal_nodes=["complete", "not_improved"],
        metadata={"description": "Project evolution workflow: measurable baseline, Sol plan, approval, normal task execution, objective comparison."},
    )
    spec.validate()
    return spec


def project_exploration_graph_v1() -> ExecutionGraphSpec:
    spec = ExecutionGraphSpec(
        name="project_exploration",
        version=1,
        graph_type="PROJECT_EXPLORATION",
        entry="prepare",
        nodes={node: GraphNodeSpec(node, node if node == "complete" else f"exploration_{node}") for node in ["prepare", "gather_project_evidence", "identify_gaps", "generate_opportunities", "normalize", "deduplicate", "score", "rank", "complete"]},
        edges=[GraphEdgeSpec("prepare", "gather_project_evidence", "success"), GraphEdgeSpec("gather_project_evidence", "identify_gaps", "success"), GraphEdgeSpec("identify_gaps", "generate_opportunities", "success"), GraphEdgeSpec("generate_opportunities", "normalize", "success"), GraphEdgeSpec("normalize", "deduplicate", "success"), GraphEdgeSpec("deduplicate", "score", "success"), GraphEdgeSpec("score", "rank", "success"), GraphEdgeSpec("rank", "complete", "success")],
        terminal_nodes=["complete"],
        metadata={"description": "Project exploration workflow: discover, deduplicate, score, rank opportunities without execution."},
    )
    spec.validate()
    return spec


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

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


class ExtensionNode(LifecycleNode):
    def __init__(self, service: ExtensionService, node_type: str) -> None:
        super().__init__(node_type)
        self.service = service

    def plan_id(self, context: GraphExecutionContext) -> str:
        return str(context.graph_run.metadata.get("extension_plan_id", ""))

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


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


class ExtensionGatherContextNode(ExtensionNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        from control_plane.projects.models import ExtensionCandidate

        candidate = ExtensionCandidate.objects.get(id=self.candidate_id(context))
        metadata = dict(context.graph_run.metadata)
        metadata["context_snapshot"] = self.service.project_context(candidate.project)
        context.graph_run.metadata = metadata
        context.graph_run.save(update_fields=["metadata", "updated_at"])
        return NodeResult("COMPLETE", "success")


class ExtensionPlanNode(ExtensionNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        from control_plane.projects.models import ExtensionCandidate

        candidate = ExtensionCandidate.objects.get(id=self.candidate_id(context))
        plan = candidate.plans.order_by("-created_at").first() or self.service.plan_with_project_brain(candidate)
        metadata = dict(context.graph_run.metadata)
        metadata["extension_plan_id"] = str(plan.id)
        context.graph_run.metadata = metadata
        context.graph_run.save(update_fields=["metadata", "updated_at"])
        return NodeResult("COMPLETE", "success", {"extension_plan_id": str(plan.id)})


class ExtensionValidatePlanNode(ExtensionNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        from control_plane.projects.models import ExtensionPlan
        from project_brain.planning import parse_project_plan_response
        import json

        plan = ExtensionPlan.objects.get(id=self.plan_id(context))
        parse_project_plan_response(json.dumps(plan.plan.get("project_plan", {})))
        return NodeResult("COMPLETE", "success")


class ApprovalNode(LifecycleNode):
    reason = "AWAITING_LIFECYCLE_APPROVAL"

    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():
            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=self.reason)
        return NodeResult("PAUSED", "awaiting", pause_reason=self.reason)


class ExtensionMaterializeNode(ExtensionNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        from control_plane.projects.models import ExtensionPlan

        plan = ExtensionPlan.objects.get(id=self.plan_id(context))
        self.service.approve_plan(plan)
        project_plan = self.service.materialize_project_dag(plan)
        return NodeResult("COMPLETE", "success", {"project_plan_id": str(project_plan.id)})


class ExtensionExecuteNode(ExtensionNode):
    destructive = True

    def __init__(self, service: ExtensionService, node_type: str, router=None, test_command: list[str] | None = None) -> None:
        super().__init__(service, node_type)
        self.router = router
        self.test_command = test_command

    def run(self, context: GraphExecutionContext) -> NodeResult:
        from control_plane.projects.models import ExtensionPlan

        if self.router is None:
            return NodeResult("COMPLETE", "success", {"executed": False})
        plan = ExtensionPlan.objects.get(id=self.plan_id(context))
        runs = self.service.execute(plan, self.router, test_command=self.test_command)
        return NodeResult("COMPLETE", "success", {"graph_run_ids": [str(run.id) for run in runs]})


class ExtensionVerifyNode(ExtensionNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        from control_plane.projects.models import ExtensionPlan

        plan = ExtensionPlan.objects.get(id=self.plan_id(context))
        verification = self.service.verify_extension(plan)
        return NodeResult("COMPLETE", "success", {"verification_id": str(verification.id), "result": verification.result})


class EvolutionNode(LifecycleNode):
    def __init__(self, service: EvolutionService, node_type: str) -> None:
        super().__init__(node_type)
        self.service = service

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

    def plan_id(self, context: GraphExecutionContext) -> str:
        return str(context.graph_run.metadata.get("evolution_plan_id", ""))


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


class EvolutionBaselineNode(EvolutionNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        from control_plane.projects.models import EvolutionCandidate

        candidate = EvolutionCandidate.objects.get(id=self.candidate_id(context))
        if not candidate.baseline_measurement:
            return NodeResult("FAILED", "failure", failure_evidence={"reason": "EVOLVE requires baseline measurement"})
        return NodeResult("COMPLETE", "success", {"baseline": candidate.baseline_measurement})


class EvolutionPlanNode(EvolutionNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        from control_plane.projects.models import EvolutionCandidate

        candidate = EvolutionCandidate.objects.get(id=self.candidate_id(context))
        plan = candidate.plans.order_by("-created_at").first() or self.service.plan_with_project_brain(candidate)
        metadata = dict(context.graph_run.metadata)
        metadata["evolution_plan_id"] = str(plan.id)
        context.graph_run.metadata = metadata
        context.graph_run.save(update_fields=["metadata", "updated_at"])
        return NodeResult("COMPLETE", "success", {"evolution_plan_id": str(plan.id)})


class EvolutionMaterializeNode(EvolutionNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        from control_plane.projects.models import EvolutionPlan

        plan = EvolutionPlan.objects.get(id=self.plan_id(context))
        self.service.approve_plan(plan)
        project_plan = self.service.materialize_work(plan)
        return NodeResult("COMPLETE", "success", {"project_plan_id": str(project_plan.id)})


class EvolutionExecuteNode(EvolutionNode):
    destructive = True

    def __init__(self, service: EvolutionService, node_type: str, router=None, test_command: list[str] | None = None) -> None:
        super().__init__(service, node_type)
        self.router = router
        self.test_command = test_command

    def run(self, context: GraphExecutionContext) -> NodeResult:
        from control_plane.projects.models import EvolutionPlan

        if self.router is None:
            return NodeResult("COMPLETE", "success", {"executed": False})
        plan = EvolutionPlan.objects.get(id=self.plan_id(context))
        runs = self.service.execute(plan, self.router, test_command=self.test_command)
        return NodeResult("COMPLETE", "success", {"graph_run_ids": [str(run.id) for run in runs]})


class EvolutionMeasureNode(EvolutionNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        from control_plane.projects.models import EvolutionPlan

        plan = EvolutionPlan.objects.get(id=self.plan_id(context))
        measurement = self.service.measure_candidate(plan)
        return NodeResult("COMPLETE", "success", {"measurement": measurement})


class EvolutionCompareNode(EvolutionNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        from control_plane.projects.models import EvolutionPlan

        plan = EvolutionPlan.objects.get(id=self.plan_id(context))
        delta = self.service.compare_baseline(plan)
        return NodeResult("COMPLETE", "success", {"delta": delta})


class EvolutionJudgeNode(EvolutionNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        from control_plane.projects.models import EvolutionPlan

        plan = EvolutionPlan.objects.get(id=self.plan_id(context))
        verification = self.service.judge_evolution(plan)
        plan.refresh_from_db()
        return NodeResult("COMPLETE", plan.verdict or "NOT_IMPROVED", {"verification_id": str(verification.id), "verdict": plan.verdict})


class ExplorationNode(LifecycleNode):
    def __init__(self, service: ExplorerService, node_type: str) -> None:
        super().__init__(node_type)
        self.service = service

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


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


class ExplorationGatherNode(ExplorationNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        from control_plane.projects.models import Exploration

        exploration = Exploration.objects.get(id=self.exploration_id(context))
        exploration.context_snapshot = self.service.project_context(exploration.project)
        exploration.save(update_fields=["context_snapshot", "updated_at"])
        return NodeResult("COMPLETE", "success")


class ExplorationGenerateNode(ExplorationNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        from control_plane.projects.models import Exploration

        exploration = Exploration.objects.get(id=self.exploration_id(context))
        opportunities = self.service.generate_opportunities(exploration)
        return NodeResult("COMPLETE", "success", {"opportunity_count": len(opportunities)})


class ExplorationRankNode(ExplorationNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        from control_plane.projects.models import Exploration

        exploration = Exploration.objects.get(id=self.exploration_id(context))
        ranked = list(exploration.opportunities.order_by("-composite_score").values_list("id", flat=True))
        exploration.metadata = {**exploration.metadata, "ranked_opportunity_ids": [str(item) for item in ranked]}
        exploration.status = "COMPLETE"
        exploration.completed_at = timezone.now()
        exploration.save(update_fields=["metadata", "status", "completed_at", "updated_at"])
        return NodeResult("COMPLETE", "success", {"ranked_opportunity_ids": [str(item) for item in ranked]})


def extension_registry(service: ExtensionService, router=None, test_command: list[str] | None = None) -> NodeHandlerRegistry:
    registry = NodeHandlerRegistry()
    for handler in [ExtensionPrepareNode(service, "extension_prepare"), ExtensionGatherContextNode(service, "extension_gather_context"), ExtensionPlanNode(service, "extension_plan_with_project_brain"), ExtensionValidatePlanNode(service, "extension_validate_plan"), ApprovalNode("extension_await_approval"), ExtensionMaterializeNode(service, "extension_materialize_project_dag"), ExtensionExecuteNode(service, "extension_execute", router, test_command), ExtensionVerifyNode(service, "extension_verify_extension")]:
        registry.register(handler)
    return registry


def evolution_registry(service: EvolutionService, router=None, test_command: list[str] | None = None) -> NodeHandlerRegistry:
    registry = NodeHandlerRegistry()
    for handler in [EvolutionSimpleNode(service, "evolution_prepare"), EvolutionBaselineNode(service, "evolution_establish_baseline"), EvolutionSimpleNode(service, "evolution_form_hypothesis"), EvolutionPlanNode(service, "evolution_plan_with_project_brain"), ApprovalNode("evolution_await_approval"), EvolutionMaterializeNode(service, "evolution_materialize_work"), EvolutionExecuteNode(service, "evolution_execute", router, test_command), EvolutionMeasureNode(service, "evolution_measure_candidate"), EvolutionCompareNode(service, "evolution_compare_baseline"), EvolutionJudgeNode(service, "evolution_judge_evolution")]:
        registry.register(handler)
    return registry


def exploration_registry(service: ExplorerService) -> NodeHandlerRegistry:
    registry = NodeHandlerRegistry()
    for handler in [ExplorationSimpleNode(service, "exploration_prepare"), ExplorationGatherNode(service, "exploration_gather_project_evidence"), ExplorationSimpleNode(service, "exploration_identify_gaps"), ExplorationGenerateNode(service, "exploration_generate_opportunities"), ExplorationSimpleNode(service, "exploration_normalize"), ExplorationSimpleNode(service, "exploration_deduplicate"), ExplorationSimpleNode(service, "exploration_score"), ExplorationRankNode(service, "exploration_rank")]:
        registry.register(handler)
    return registry