web_ui/frontend/src/components/graph/GraphEditor.tsx

import { useCallback, useMemo, useState } from "react";
import {
  ReactFlow,
  Background,
  Controls,
  MiniMap,
  type Edge,
  type NodeTypes,
  type EdgeTypes,
} from "@xyflow/react";
import "@xyflow/react/dist/style.css";

import { AgentNode } from "./AgentNode";
import { ConditionalEdge } from "./ConditionalEdge";
import { EdgeConditionDialog } from "./EdgeConditionDialog";
import { AgentForm } from "@/components/agents/AgentForm";
import { useGraphStore, type AgentNodeData, type EdgeData } from "@/stores/graphStore";
import { useAgentStore } from "@/stores/agentStore";
import type { AgentCreateRequest } from "@/types/agent";

export function GraphEditor() {
  const {
    nodes,
    edges,
    onNodesChange,
    onEdgesChange,
    onConnect,
    addAgentNode,
    removeNode,
    setEdgeCondition,
    editingNodeId,
    setEditingNodeId,
    updateNodeData,
  } = useGraphStore();

  const agents = useAgentStore((s) => s.agents);
  const updateAgent = useAgentStore((s) => s.updateAgent);

  const [selectedEdge, setSelectedEdge] = useState<Edge<EdgeData> | null>(null);
  const [conditionDialogOpen, setConditionDialogOpen] = useState(false);

  const nodeTypes: NodeTypes = useMemo(() => ({ agentNode: AgentNode }), []);
  const edgeTypes: EdgeTypes = useMemo(() => ({ conditionEdge: ConditionalEdge }), []);

  const handleEdgeClick = useCallback((_: React.MouseEvent, edge: Edge) => {
    setSelectedEdge(edge as Edge<EdgeData>);
    setConditionDialogOpen(true);
  }, []);

  const handleEdgeConditionSave = useCallback(
    (edgeId: string, condition: string, weight: number) => {
      setEdgeCondition(edgeId, condition, weight);
    },
    [setEdgeCondition]
  );

  const handleDragOver = useCallback((e: React.DragEvent) => {
    e.preventDefault();
    e.dataTransfer.dropEffect = "move";
  }, []);

  const handleDrop = useCallback(
    (e: React.DragEvent) => {
      e.preventDefault();
      const agentId = e.dataTransfer.getData("application/agent-id");
      if (!agentId) return;

      const agent = agents.find((a) => a.agent_id === agentId);
      if (!agent) return;

      const reactFlowBounds = (e.target as HTMLElement)
        .closest(".react-flow")
        ?.getBoundingClientRect();
      if (!reactFlowBounds) return;

      const position = {
        x: e.clientX - reactFlowBounds.left - 100,
        y: e.clientY - reactFlowBounds.top - 60,
      };

      addAgentNode(agent, position);
    },
    [agents, addAgentNode]
  );

  const handleKeyDown = useCallback(
    (e: React.KeyboardEvent) => {
      if (e.key === "Delete" || e.key === "Backspace") {
        const selectedNodes = nodes.filter((n) => n.selected);
        selectedNodes.forEach((n) => removeNode(n.id));
      }
    },
    [nodes, removeNode]
  );

  // Build AgentProfile-like object from the node being edited
  const editingNode = editingNodeId
    ? nodes.find((n) => n.id === editingNodeId)
    : null;

  const editingAgentProfile = editingNode
    ? {
        agent_id: editingNode.data.agentId,
        display_name: editingNode.data.displayName,
        persona: editingNode.data.persona,
        description: editingNode.data.description,
        llm_backbone: editingNode.data.llmBackbone,
        tools: editingNode.data.tools,
        llm_config: null,
        input_schema: null,
        output_schema: null,
      }
    : null;

  const handleAgentEditSubmit = async (data: AgentCreateRequest) => {
    if (!editingNodeId) return;
    // Update the node on the canvas
    updateNodeData(editingNodeId, {
      displayName: data.display_name,
      persona: data.persona || "",
      description: data.description || "",
      llmBackbone: data.llm_backbone || null,
      tools: data.tools || [],
    });
    // Also update the agent on the backend
    try {
      await updateAgent(editingNodeId, data);
    } catch {
      // agent may not exist on server yet (template-only)
    }
  };

  return (
    <div className="flex-1 h-full" onKeyDown={handleKeyDown} tabIndex={0}>
      <ReactFlow
        nodes={nodes}
        edges={edges}
        onNodesChange={onNodesChange}
        onEdgesChange={onEdgesChange}
        onConnect={onConnect}
        onEdgeClick={handleEdgeClick}
        onDragOver={handleDragOver}
        onDrop={handleDrop}
        nodeTypes={nodeTypes}
        edgeTypes={edgeTypes}
        fitView
        snapToGrid
        snapGrid={[15, 15]}
        defaultEdgeOptions={{
          type: "conditionEdge",
          animated: true,
        }}
        deleteKeyCode={["Backspace", "Delete"]}
      >
        <Background gap={15} size={1} />
        <Controls />
        <MiniMap
          nodeColor={(n) => {
            const data = n.data as unknown as AgentNodeData;
            switch (data.executionStatus) {
              case "running":
                return "#3b82f6";
              case "completed":
                return "#22c55e";
              case "error":
                return "#ef4444";
              case "pending":
                return "#eab308";
              default:
                return "#94a3b8";
            }
          }}
          maskColor="rgba(0, 0, 0, 0.1)"
        />
      </ReactFlow>

      <EdgeConditionDialog
        open={conditionDialogOpen}
        onOpenChange={setConditionDialogOpen}
        edgeId={selectedEdge?.id || ""}
        currentCondition={(selectedEdge?.data as EdgeData | undefined)?.condition}
        currentWeight={(selectedEdge?.data as EdgeData | undefined)?.weight ?? 1.0}
        onSave={handleEdgeConditionSave}
      />

      <AgentForm
        open={!!editingNodeId}
        onOpenChange={(open) => { if (!open) setEditingNodeId(null); }}
        agent={editingAgentProfile}
        onSubmit={handleAgentEditSubmit}
      />
    </div>
  );
}