import { useCallback, useMemo, memo } 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 { useShallow } from 'zustand/react/shallow';
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" />,
  vm_lxc: <Box className="w-4 h-4" />,
  vm_qemu: <Server className="w-5 h-5" />,
  systemd_service: <Box className="w-4 h-4" />,
  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" />,
};

const CustomNode = memo(function CustomNode({ data, selected, id }: NodeProps) {
  const highlightPath = useTopologyStore((s) => s.highlightPath);
  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');
  const isHighlighted = highlightPath.includes(id);
  const isDimmed = highlightPath.length > 0 && !isHighlighted;

  return (
    <div
      className={`px-4 py-3 rounded-xl border-2 transition-all duration-200 ${selected
          ? 'border-sky-400 shadow-lg shadow-sky-400/20 scale-[1.02]'
          : isHighlighted
            ? 'border-indigo-400 shadow-lg shadow-indigo-400/20'
            : 'border-slate-600 hover:border-slate-500 hover:shadow-md hover:shadow-slate-700/30 hover:scale-[1.01]'
        }`}
      style={{
        backgroundColor: isDimmed ? '#0F172A' : '#1E293B',
        minWidth: '140px',
        opacity: isDimmed ? 0.4 : 1
      }}
      role="treeitem"
      aria-label={`${nodeData.label || 'Node'}, ${nodeData.type?.replace(/_/g, ' ') || 'unknown type'}, ${nodeData.status || 'unknown status'}`}
    >
      <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 }} aria-hidden="true">
            {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 }}
          aria-label={`Status: ${nodeData.status || 'unknown'}`}
        />
      </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,
    highlightPath
  } = useTopologyStore(useShallow((s) => ({
    edges: s.edges,
    selectedNodeId: s.selectedNodeId,
    setSelectedNode: s.setSelectedNode,
    getFilteredNodes: s.getFilteredNodes,
    orientation: s.orientation,
    viewMode: s.viewMode,
    highlightPath: s.highlightPath,
  })));

  // Memoize the layout computation instead of useState + useEffect
  const { nodes, edges } = useMemo(() => {
    const filteredNodes = getFilteredNodes();

    if (filteredNodes.length === 0) {
      return { nodes: [] as Node[], edges: [] as Edge[] };
    }

    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 => {
        const isPathEdge = highlightPath.includes(edge.source) && highlightPath.includes(edge.target);
        const isSelected = edge.source === selectedNodeId || edge.target === selectedNodeId;
        return {
          id: edge.id,
          source: edge.source,
          target: edge.target,
          type: 'smoothstep',
          animated: isSelected || isPathEdge,
          style: {
            stroke: isSelected
              ? '#38BDF8'
              : isPathEdge
                ? '#818CF8'
                : '#475569',
            strokeWidth: isSelected || isPathEdge
              ? 2
              : 1,
            opacity: highlightPath.length > 0 && !isPathEdge ? 0.3 : 1
          },
          markerEnd: {
            type: 'arrowclosed' as const,
            color: isSelected
              ? '#38BDF8'
              : isPathEdge
                ? '#818CF8'
                : '#475569',
          },
        };
      });

    return getLayoutedElements(newNodes, newEdges, orientation);
  }, [getFilteredNodes, storeEdges, selectedNodeId, orientation, viewMode, highlightPath]);

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

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

  return (
    <div className="w-full h-full" role="application" aria-label="Network topology graph">
      <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>
  );
}