homelab-topology/src/components/Graph/TopologyGraph.tsx

import { useCallback, useEffect, useState } from 'react';
import {
  ReactFlow,
  Background,
  Controls,
  MiniMap,
  Node,
  Edge,
  NodeProps,
  Handle,
  Position,
} from '@xyflow/react';
import '@xyflow/react/dist/style.css';
import dagre from 'dagre';
import { useTopologyStore } from '../../store/topologyStore';
import { getNodeColor, getStatusColor } from '../../utils/colors';
import { Server, Network, Wifi, Box, Database, Folder } from 'lucide-react';
import { NodeType, ServiceCategory } from '../../types';

const nodeIcons: Record<NodeType, React.ReactNode> = {
  gateway: <Network className="w-5 h-5" />,
  vlan: <Network className="w-4 h-4" />,
  wifi: <Wifi className="w-4 h-4" />,
  host_physical: <Server className="w-5 h-5" />,
  host_vm: <Server className="w-5 h-5" />,
  host_container: <Server className="w-5 h-5" />,
  service: <Box className="w-4 h-4" />,
  volume: <Database className="w-4 h-4" />,
  mount: <Database className="w-4 h-4" />,
  path: <Folder className="w-4 h-4" />,
};

function CustomNode({ data, selected }: NodeProps) {
  const nodeData = data as { type?: NodeType; category?: ServiceCategory; status?: 'running' | 'stopped' | 'unknown'; label?: string; ip?: string };
  const nodeColor = getNodeColor(nodeData.type || 'service', nodeData.category);
  const statusColor = getStatusColor(nodeData.status || 'unknown');

  return (
    <div
      className={`px-4 py-3 rounded-xl border-2 transition-all ${
        selected 
          ? 'border-sky-400 shadow-lg shadow-sky-400/20' 
          : 'border-slate-600 hover:border-slate-500'
      }`}
      style={{ 
        backgroundColor: '#1E293B',
        minWidth: '140px'
      }}
    >
      <Handle type="target" position={Position.Left} className="!bg-slate-400" />
      
      <div className="flex items-center gap-3">
        <div 
          className="w-10 h-10 rounded-lg flex items-center justify-center"
          style={{ backgroundColor: `${nodeColor}20` }}
        >
          <div style={{ color: nodeColor }}>
            {nodeIcons[nodeData.type || 'service']}
          </div>
        </div>
        
        <div className="flex-1 min-w-0">
          <div className="text-sm font-medium text-white truncate">
            {nodeData.label}
          </div>
          {nodeData.ip && (
            <div className="text-xs text-slate-500 font-mono">
              {nodeData.ip}
            </div>
          )}
        </div>
        
        <div 
          className="w-2.5 h-2.5 rounded-full"
          style={{ backgroundColor: statusColor }}
        />
      </div>
      
      <Handle type="source" position={Position.Right} className="!bg-slate-400" />
    </div>
  );
}

const nodeTypes = {
  custom: CustomNode,
};

const nodeWidth = 180;
const nodeHeight = 70;

function getLayoutedElements(nodes: Node[], edges: Edge[], direction: 'LR' | 'TB') {
  if (nodes.length === 0) return { nodes: [], edges: [] };
  
  const dagreGraph = new dagre.graphlib.Graph();
  dagreGraph.setDefaultEdgeLabel(() => ({}));
  dagreGraph.setGraph({ rankdir: direction, nodesep: 50, ranksep: 100 });

  nodes.forEach((node) => {
    dagreGraph.setNode(node.id, { width: nodeWidth, height: nodeHeight });
  });

  edges.forEach((edge) => {
    dagreGraph.setEdge(edge.source, edge.target);
  });

  dagre.layout(dagreGraph);

  const layoutedNodes = nodes.map((node) => {
    const nodeWithPosition = dagreGraph.node(node.id);
    if (!nodeWithPosition) return node;
    return {
      ...node,
      position: {
        x: nodeWithPosition.x - nodeWidth / 2,
        y: nodeWithPosition.y - nodeHeight / 2,
      },
    };
  });

  return { nodes: layoutedNodes, edges };
}

export default function TopologyGraph() {
  const { 
    edges: storeEdges, 
    selectedNodeId, 
    setSelectedNode,
    getFilteredNodes,
    orientation,
    viewMode 
  } = useTopologyStore();

  const [nodes, setNodes] = useState<Node[]>([]);
  const [edges, setEdges] = useState<Edge[]>([]);

  useEffect(() => {
    const filteredNodes = getFilteredNodes();
    
    if (filteredNodes.length === 0) {
      setNodes([]);
      setEdges([]);
      return;
    }

    const nodeIds = new Set(filteredNodes.map(n => n.id));
    
    const newNodes: Node[] = filteredNodes.map(node => ({
      id: node.id,
      type: 'custom',
      position: { x: 0, y: 0 },
      data: {
        label: node.name,
        type: node.type,
        status: node.data.status,
        category: node.data.category,
        ip: node.data.ip,
      },
      selected: node.id === selectedNodeId,
    }));

    const newEdges: Edge[] = storeEdges
      .filter(e => nodeIds.has(e.source) && nodeIds.has(e.target))
      .map(edge => ({
        id: edge.id,
        source: edge.source,
        target: edge.target,
        type: 'smoothstep',
        animated: edge.source === selectedNodeId || edge.target === selectedNodeId,
        style: { 
          stroke: edge.source === selectedNodeId || edge.target === selectedNodeId 
            ? '#38BDF8' 
            : '#475569',
          strokeWidth: edge.source === selectedNodeId || edge.target === selectedNodeId 
            ? 2 
            : 1,
        },
        markerEnd: {
          type: 'arrowclosed' as const,
          color: edge.source === selectedNodeId || edge.target === selectedNodeId 
            ? '#38BDF8' 
            : '#475569',
        },
      }));

    const { nodes: layoutedNodes, edges: layoutedEdges } = getLayoutedElements(
      newNodes,
      newEdges,
      orientation
    );

    setNodes(layoutedNodes);
    setEdges(layoutedEdges);
  }, [getFilteredNodes, storeEdges, selectedNodeId, orientation, viewMode]);

  const onNodeClick = useCallback((_: React.MouseEvent, node: Node) => {
    setSelectedNode(node.id);
  }, [setSelectedNode]);

  const onPaneClick = useCallback(() => {
    setSelectedNode(null);
  }, [setSelectedNode]);

  return (
    <div className="w-full h-full">
      <ReactFlow
        nodes={nodes}
        edges={edges}
        onNodeClick={onNodeClick}
        onPaneClick={onPaneClick}
        nodeTypes={nodeTypes}
        fitView
        fitViewOptions={{ padding: 0.2 }}
        minZoom={0.1}
        maxZoom={2}
        defaultEdgeOptions={{
          type: 'smoothstep',
        }}
        proOptions={{ hideAttribution: true }}
      >
        <Background color="#334155" gap={20} size={1} />
        <Controls className="!bg-slate-700 !border-slate-600 !rounded-lg !shadow-lg" />
        <MiniMap
          className="!bg-slate-800 !border-slate-700"
          nodeColor={(node) => getNodeColor(node.data?.type as NodeType || 'service', (node.data?.category as ServiceCategory) || undefined)}
          maskColor="rgba(15, 23, 42, 0.8)"
        />
      </ReactFlow>
    </div>
  );
}