OpenFlare/frontend/app/(main)/waf/rules/editor/components/rule-flow-canvas.tsx

'use client';

import {useCallback, useEffect, useRef} from 'react';
import {
  addEdge,
  applyEdgeChanges,
  Background,
  Controls,
  MiniMap,
  ReactFlow,
  useEdgesState,
  useNodesState,
  type Connection,
  type Edge,
  type EdgeChange,
  type Node,
  type NodeChange,
  type ReactFlowInstance,
} from '@xyflow/react';
import {Trash2} from 'lucide-react';
import '@xyflow/react/dist/style.css';

import {Button} from '@/components/ui/button';
import type {WAFRuleEdge, WAFRuleGraph, WAFRuleNode} from '@/lib/services/openflare';

import {
  acceptedNodeChanges,
  filterRemovableNodeIds,
  type GraphErrorTarget,
  isConnectionAllowed,
  isPersistentEdgeChange,
} from './editor-behavior';
import {type GraphIssue, removeEdgeFromGraph, removeNodeFromGraph} from './graph-validation';
import {NodeLibrary} from './node-library';
import {RuleNode, type RuleFlowNodeData} from './rule-node';

const nodeTypes = {rule: RuleNode};

type FlowNode = Node<RuleFlowNodeData>;
type FlowEdge = Edge;

interface RuleFlowCanvasProps {
  graph: WAFRuleGraph;
  issues: GraphIssue[];
  selectedId?: string;
  selectedEdgeId?: string;
  focusTarget?: GraphErrorTarget;
  onGraphChange: (graph: WAFRuleGraph, persistent?: boolean) => void;
  onSelect: (id?: string) => void;
  onSelectEdge: (id?: string) => void;
}

export function RuleFlowCanvas({
  graph,
  issues,
  selectedId,
  selectedEdgeId,
  focusTarget,
  onGraphChange,
  onSelect,
  onSelectEdge,
}: RuleFlowCanvasProps) {
  const instance = useRef<ReactFlowInstance<FlowNode, FlowEdge> | null>(null);
  const [nodes, setNodes, applyFlowNodeChanges] = useNodesState<FlowNode>(buildFlowNodes(graph, issues, selectedId));
  const [edges, setEdges, applyFlowEdgeChanges] = useEdgesState<FlowEdge>(buildFlowEdges(graph, selectedId, selectedEdgeId));

  useEffect(() => {
    setNodes((current) => buildFlowNodes(graph, issues, selectedId, current));
  }, [graph, issues, selectedId, setNodes]);

  useEffect(() => {
    setEdges((current) => buildFlowEdges(graph, selectedId, selectedEdgeId, current));
  }, [graph, selectedEdgeId, selectedId, setEdges]);

  useEffect(() => {
    if (!focusTarget || !instance.current) return;
    if (focusTarget.kind === 'node') {
      void instance.current.fitView({nodes: [{id: focusTarget.id}], duration: 350, maxZoom: 1.4});
      return;
    }
    const edge = graph.edges.find((item) => item.id === focusTarget.id);
    if (edge) void instance.current.fitView({nodes: [{id: edge.source}, {id: edge.target}], duration: 350, maxZoom: 1.4});
  }, [focusTarget, graph.edges]);

  const onNodesChange = useCallback((changes: NodeChange<FlowNode>[]) => {
    const accepted = acceptedNodeChanges(graph.nodes, changes);
    if (accepted.changes.length === 0) return;

    applyFlowNodeChanges(accepted.changes);
    if (!accepted.persistent) return;

    const removals = filterRemovableNodeIds(
      graph.nodes,
      accepted.changes.filter((change) => change.type === 'remove').map((change) => change.id),
    );
    let next = graph;
    for (const id of removals) next = removeNodeFromGraph(next, id);

    const positions = new Map<string, {x: number; y: number}>();
    for (const change of accepted.changes) {
      if (change.type === 'position' && change.dragging === false && change.position) {
        positions.set(change.id, change.position);
      }
    }
    next = {...next, nodes: next.nodes.map((node) => ({...node, position: positions.get(node.id) ?? node.position}))};
    onGraphChange(next);

    if (selectedId && removals.includes(selectedId)) onSelect(undefined);
  }, [applyFlowNodeChanges, graph, onGraphChange, onSelect, selectedId]);

  const onEdgesChange = useCallback((changes: EdgeChange<FlowEdge>[]) => {
    applyFlowEdgeChanges(changes);
    const persistent = changes.filter(isPersistentEdgeChange);
    if (persistent.length === 0) return;

    const next = applyEdgeChanges(persistent, edges);
    onGraphChange({...graph, edges: next.map(toRuleEdge)});
    if (selectedEdgeId && persistent.some((change) => change.type === 'remove' && change.id === selectedEdgeId)) {
      onSelectEdge(undefined);
    }
  }, [applyFlowEdgeChanges, edges, graph, onGraphChange, onSelectEdge, selectedEdgeId]);

  const isValidConnection = useCallback((connection: Edge | Connection) => {
    return isConnectionAllowed(graph, connection);
  }, [graph]);

  const onConnect = useCallback((connection: Connection) => {
    if (!isValidConnection(connection)) return;
    const next = addEdge({...connection, id: `${connection.source}-${connection.sourceHandle}-${connection.target}`}, edges);
    setEdges(next);
    onGraphChange({...graph, edges: next.map(toRuleEdge)});
  }, [edges, graph, isValidConnection, onGraphChange, setEdges]);

  const addNode = useCallback((type: 'ip_match' | 'geo_match' | 'pow' | 'block') => {
    const id = `${type}-${crypto.randomUUID().slice(0, 8)}`;
    const config = type === 'ip_match'
      ? {ips: [], cidrs: [], ip_group_ids: []}
      : type === 'geo_match'
        ? {countries: [], regions: []}
        : type === 'pow'
          ? {algorithm: 'fast' as const, difficulty: 4, session_ttl: 3600, challenge_ttl: 300}
          : {status_code: 403, response_body: ''};
    onGraphChange({...graph, nodes: [...graph.nodes, {id, type, position: {x: 240, y: 140 + graph.nodes.length * 24}, config} as WAFRuleNode]});
    onSelectEdge(undefined);
    onSelect(id);
  }, [graph, onGraphChange, onSelect, onSelectEdge]);

  const selectedNode = graph.nodes.find((node) => node.id === selectedId);
  const canDeleteNode = Boolean(selectedNode && !['start', 'allow'].includes(selectedNode.type));
  const canDeleteEdge = Boolean(selectedEdgeId && graph.edges.some((edge) => edge.id === selectedEdgeId));
  const canDelete = canDeleteNode || canDeleteEdge;
  const deleteLabel = canDeleteEdge ? '删除连线' : canDeleteNode ? '删除节点' : selectedNode ? '系统节点不可删除' : '请选择可删除项';

  const deleteSelection = useCallback(() => {
    if (canDeleteEdge && selectedEdgeId) {
      setEdges((current) => current.filter((edge) => edge.id !== selectedEdgeId));
      onGraphChange(removeEdgeFromGraph(graph, selectedEdgeId));
      onSelectEdge(undefined);
      return;
    }
    if (canDeleteNode && selectedId) {
      setNodes((current) => current.filter((node) => node.id !== selectedId));
      setEdges((current) => current.filter((edge) => edge.source !== selectedId && edge.target !== selectedId));
      onGraphChange(removeNodeFromGraph(graph, selectedId));
      onSelect(undefined);
    }
  }, [canDeleteEdge, canDeleteNode, graph, onGraphChange, onSelect, onSelectEdge, selectedEdgeId, selectedId, setEdges, setNodes]);

  return (
    <section className="relative min-w-0 flex-1 bg-muted/20">
      <div className="absolute left-4 top-4 z-10 rounded-lg border bg-background/95 p-2 shadow-sm backdrop-blur">
        <NodeLibrary onAdd={addNode}/>
      </div>
      <div className="absolute right-4 top-4 z-10">
        <Button variant="destructive" size="sm" disabled={!canDelete} onClick={deleteSelection}>
          <Trash2 data-icon="inline-start"/>
          {deleteLabel}
        </Button>
      </div>
      <ReactFlow
        nodes={nodes}
        edges={edges}
        nodeTypes={nodeTypes}
        onInit={(value) => { instance.current = value; }}
        onNodesChange={onNodesChange}
        onEdgesChange={onEdgesChange}
        onConnect={onConnect}
        isValidConnection={isValidConnection}
        onBeforeDelete={async ({nodes: deletingNodes}) => !deletingNodes.some((node) => ['start', 'allow'].includes(node.data.rule.type))}
        onNodeClick={(_, node) => { onSelectEdge(undefined); onSelect(node.id); }}
        onEdgeClick={(_, edge) => { onSelect(undefined); onSelectEdge(edge.id); }}
        onPaneClick={() => { onSelect(undefined); onSelectEdge(undefined); }}
        fitView
        deleteKeyCode={['Backspace', 'Delete']}
      >
        <Background gap={20} size={1}/>
        <MiniMap pannable zoomable/>
        <Controls/>
      </ReactFlow>
    </section>
  );
}

function buildFlowNodes(graph: WAFRuleGraph, issues: GraphIssue[], selectedId?: string, current: FlowNode[] = []): FlowNode[] {
  const existing = new Map(current.map((node) => [node.id, node]));
  const issueCounts = new Map<string, number>();
  for (const issue of issues) {
    if (issue.nodeId) issueCounts.set(issue.nodeId, (issueCounts.get(issue.nodeId) ?? 0) + 1);
  }
  return graph.nodes.map((rule) => ({
    ...existing.get(rule.id),
    id: rule.id,
    type: 'rule',
    position: rule.position,
    selected: rule.id === selectedId,
    data: {rule, issues: issueCounts.get(rule.id) ?? 0},
  }));
}

function buildFlowEdges(graph: WAFRuleGraph, selectedId?: string, selectedEdgeId?: string, current: FlowEdge[] = []): FlowEdge[] {
  const existing = new Map(current.map((edge) => [edge.id, edge]));
  return graph.edges.map((edge) => ({
    ...existing.get(edge.id),
    id: edge.id,
    source: edge.source,
    sourceHandle: edge.source_handle,
    target: edge.target,
    selected: edge.id === selectedEdgeId,
    animated: selectedId === edge.source || selectedId === edge.target,
  }));
}

function toRuleEdge(edge: Edge): WAFRuleEdge {
  return {id: edge.id, source: edge.source, source_handle: edge.sourceHandle ?? '', target: edge.target};
}