from __future__ import annotations

from agents.control_room import AgentControlRoomService
from control_plane.agents.models import AgentVersion
from graph.native_runtime import GraphExecutionContext
from graph.registry import NodeHandlerRegistry, NodeResult
from graph.spec import ExecutionGraphSpec, GraphEdgeSpec, GraphNodeSpec


def agent_investigation_graph_v1() -> ExecutionGraphSpec:
    nodes = ["prepare", "gather_agent_evidence", "smart_investigation", "improvement_candidates", "complete"]
    spec = ExecutionGraphSpec(
        name="agent_investigation",
        version=1,
        graph_type="AGENT_INVESTIGATION",
        entry="prepare",
        nodes={node: GraphNodeSpec(node, node if node == "complete" else f"agent_{node}") for node in nodes},
        edges=[GraphEdgeSpec(nodes[index], nodes[index + 1], "success") for index in range(len(nodes) - 1)],
        terminal_nodes=["complete"],
        metadata={"description": "Agent investigation workflow using Control Room telemetry and Progeny Smart Investigation."},
    )
    spec.validate()
    return spec


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

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

    def version(self, context: GraphExecutionContext) -> AgentVersion:
        return AgentVersion.objects.get(id=context.graph_run.metadata["agent_version_id"])


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


class AgentGatherEvidenceNode(AgentInvestigationNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        version = self.version(context)
        evidence = {"performance": self.service.get_agent_performance(version.id), "health": self.service.get_agent_health(version.id), "usage": self.service.get_agent_usage(version.id)}
        metadata = dict(context.graph_run.metadata)
        metadata["agent_evidence"] = evidence
        context.graph_run.metadata = metadata
        context.graph_run.save(update_fields=["metadata", "updated_at"])
        return NodeResult("COMPLETE", "success", evidence)


class AgentSmartInvestigationNode(AgentInvestigationNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        investigation = self.service.investigate_agent(self.version(context))
        metadata = dict(context.graph_run.metadata)
        metadata["progeny_investigation_id"] = str(investigation.id)
        context.graph_run.metadata = metadata
        context.graph_run.save(update_fields=["metadata", "updated_at"])
        return NodeResult("COMPLETE", "success", {"investigation_id": str(investigation.id)})


class AgentImprovementCandidateNode(AgentInvestigationNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        from control_plane.agents.models import ProgenyInvestigation

        investigation_id = context.graph_run.metadata.get("progeny_investigation_id")
        investigation = ProgenyInvestigation.objects.get(id=investigation_id)
        candidate = self.service.create_improvement_candidate(investigation)
        return NodeResult("COMPLETE", "success", {"improvement_candidate_id": str(candidate.id)})


def agent_investigation_registry(service: AgentControlRoomService) -> NodeHandlerRegistry:
    registry = NodeHandlerRegistry()
    for handler in [AgentSimpleNode(service, "agent_prepare"), AgentGatherEvidenceNode(service, "agent_gather_agent_evidence"), AgentSmartInvestigationNode(service, "agent_smart_investigation"), AgentImprovementCandidateNode(service, "agent_improvement_candidates")]:
        registry.register(handler)
    return registry