stack/apps/web/src/components/mindmap/ReactFlowEditor.tsx

"use client";

import { useCallback, useEffect, useMemo, useState } from "react";
import type { Connection, Node, Edge, NodeTypes } from "@xyflow/react";
import {
  ReactFlow,
  Background,
  Controls,
  MiniMap,
  Panel,
  useNodesState,
  useEdgesState,
  addEdge,
  MarkerType,
  BackgroundVariant,
} from "@xyflow/react";
import "@xyflow/react/dist/style.css";

import type {
  GraphData,
  KnowledgeNode,
  KnowledgeEdge,
  EdgeCreateInput,
} from "./hooks/useGraphData";
import { ConceptNode } from "./nodes/ConceptNode";
import { TaskNode } from "./nodes/TaskNode";
import { IdeaNode } from "./nodes/IdeaNode";
import { ProjectNode } from "./nodes/ProjectNode";

// Node type to color mapping
const NODE_COLORS: Record<string, string> = {
  concept: "#6366f1", // indigo
  idea: "#f59e0b", // amber
  task: "#10b981", // emerald
  project: "#3b82f6", // blue
  person: "#ec4899", // pink
  note: "#8b5cf6", // violet
  question: "#f97316", // orange
};

// Relation type to label mapping
const RELATION_LABELS: Record<string, string> = {
  relates_to: "relates to",
  part_of: "part of",
  depends_on: "depends on",
  mentions: "mentions",
  blocks: "blocks",
  similar_to: "similar to",
  derived_from: "derived from",
};

interface ReactFlowEditorProps {
  graphData: GraphData;
  onNodeSelect?: (node: KnowledgeNode | null) => void;
  onNodeUpdate?: (id: string, updates: Partial<KnowledgeNode>) => void;
  onNodeDelete?: (id: string) => void;
  onEdgeCreate?: (edge: EdgeCreateInput) => void;
  className?: string;
  readOnly?: boolean;
}

// Custom node types
const nodeTypes: NodeTypes = {
  concept: ConceptNode,
  task: TaskNode,
  idea: IdeaNode,
  project: ProjectNode,
  default: ConceptNode,
};

function convertToReactFlowNodes(nodes: KnowledgeNode[]): Node[] {
  // Simple grid layout for initial positioning
  const COLS = 4;
  const X_SPACING = 250;
  const Y_SPACING = 150;

  return nodes.map((node, index) => ({
    id: node.id,
    type: node.node_type in nodeTypes ? node.node_type : "default",
    position: {
      x: (index % COLS) * X_SPACING + Math.random() * 50,
      y: Math.floor(index / COLS) * Y_SPACING + Math.random() * 30,
    },
    data: {
      label: node.title,
      content: node.content,
      nodeType: node.node_type,
      tags: node.tags,
      domain: node.domain,
      id: node.id,
      metadata: node.metadata,
      created_at: node.created_at,
      updated_at: node.updated_at,
    },
    style: {
      borderColor: NODE_COLORS[node.node_type] || NODE_COLORS.concept,
    },
  }));
}

function convertToReactFlowEdges(edges: KnowledgeEdge[]): Edge[] {
  return edges.map((edge) => ({
    // Use stable ID based on source, target, and relation type
    id: `${edge.source_id}-${edge.target_id}-${edge.relation_type}`,
    source: edge.source_id,
    target: edge.target_id,
    label: RELATION_LABELS[edge.relation_type] || edge.relation_type,
    type: "smoothstep",
    animated: edge.relation_type === "depends_on" || edge.relation_type === "blocks",
    markerEnd: {
      type: MarkerType.ArrowClosed,
      width: 20,
      height: 20,
    },
    data: {
      relationType: edge.relation_type,
      weight: edge.weight,
    },
    style: {
      strokeWidth: Math.max(1, edge.weight * 3),
      opacity: 0.6 + edge.weight * 0.4,
    },
  }));
}

export function ReactFlowEditor({
  graphData,
  onNodeSelect,
  onNodeUpdate,
  onNodeDelete,
  onEdgeCreate,
  className = "",
  readOnly = false,
}: ReactFlowEditorProps) {
  const [selectedNode, setSelectedNode] = useState<string | null>(null);

  const initialNodes = useMemo(() => convertToReactFlowNodes(graphData.nodes), [graphData.nodes]);

  const initialEdges = useMemo(() => convertToReactFlowEdges(graphData.edges), [graphData.edges]);

  const [nodes, setNodes, onNodesChange] = useNodesState(initialNodes);
  const [edges, setEdges, onEdgesChange] = useEdgesState(initialEdges);

  // Update nodes/edges when graphData changes
  useEffect(() => {
    setNodes(convertToReactFlowNodes(graphData.nodes));
    setEdges(convertToReactFlowEdges(graphData.edges));
  }, [graphData, setNodes, setEdges]);

  const onConnect = useCallback(
    (params: Connection) => {
      if (readOnly || !params.source || !params.target) return;

      // Create edge in backend
      if (onEdgeCreate) {
        onEdgeCreate({
          source_id: params.source,
          target_id: params.target,
          relation_type: "relates_to",
          weight: 1.0,
          metadata: {},
        });
      }

      setEdges((eds) =>
        addEdge(
          {
            ...params,
            type: "smoothstep",
            markerEnd: { type: MarkerType.ArrowClosed },
          },
          eds
        )
      );
    },
    [readOnly, onEdgeCreate, setEdges]
  );

  const onNodeClick = useCallback(
    (_event: React.MouseEvent, node: Node) => {
      setSelectedNode(node.id);
      if (onNodeSelect) {
        const knowledgeNode = graphData.nodes.find((n) => n.id === node.id);
        onNodeSelect(knowledgeNode || null);
      }
    },
    [graphData.nodes, onNodeSelect]
  );

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

  const onNodeDragStop = useCallback(
    (_event: React.MouseEvent, node: Node) => {
      if (readOnly) return;
      // Could save position to metadata if needed
      if (onNodeUpdate) {
        onNodeUpdate(node.id, {
          metadata: { position: node.position },
        });
      }
    },
    [readOnly, onNodeUpdate]
  );

  const handleDeleteSelected = useCallback(() => {
    if (readOnly || !selectedNode) return;

    if (onNodeDelete) {
      onNodeDelete(selectedNode);
    }

    setNodes((nds) => nds.filter((n) => n.id !== selectedNode));
    setEdges((eds) => eds.filter((e) => e.source !== selectedNode && e.target !== selectedNode));
    setSelectedNode(null);
  }, [readOnly, selectedNode, onNodeDelete, setNodes, setEdges]);

  // Keyboard shortcuts
  useEffect(() => {
    const handleKeyDown = (event: KeyboardEvent) => {
      if (readOnly) return;

      if (event.key === "Delete" || event.key === "Backspace") {
        if (selectedNode && document.activeElement === document.body) {
          event.preventDefault();
          handleDeleteSelected();
        }
      }
    };

    document.addEventListener("keydown", handleKeyDown);
    return (): void => {
      document.removeEventListener("keydown", handleKeyDown);
    };
  }, [readOnly, selectedNode, handleDeleteSelected]);

  const isDark =
    typeof window !== "undefined" && document.documentElement.classList.contains("dark");

  return (
    <div className={`w-full h-full ${className}`} style={{ minHeight: "500px" }}>
      <ReactFlow
        nodes={nodes}
        edges={edges}
        {...(!readOnly && {
          onNodesChange,
          onEdgesChange,
        })}
        onConnect={onConnect}
        onNodeClick={onNodeClick}
        onPaneClick={onPaneClick}
        onNodeDragStop={onNodeDragStop}
        nodeTypes={nodeTypes}
        fitView
        attributionPosition="bottom-left"
        proOptions={{ hideAttribution: true }}
        className="bg-gray-50 dark:bg-gray-900"
      >
        <Background
          variant={BackgroundVariant.Dots}
          gap={20}
          size={1}
          color={isDark ? "#374151" : "#e5e7eb"}
        />
        <Controls
          showZoom
          showFitView
          showInteractive={!readOnly}
          className="bg-white dark:bg-gray-800 border-gray-200 dark:border-gray-700"
        />
        <MiniMap
          nodeColor={(node) => NODE_COLORS[node.data?.nodeType as string] || "#6366f1"}
          maskColor={isDark ? "rgba(0, 0, 0, 0.8)" : "rgba(255, 255, 255, 0.8)"}
          className="bg-white dark:bg-gray-800 border-gray-200 dark:border-gray-700"
        />
        <Panel position="top-right" className="bg-white dark:bg-gray-800 p-2 rounded shadow">
          <div className="text-sm text-gray-600 dark:text-gray-400">
            {graphData.nodes.length} nodes, {graphData.edges.length} edges
          </div>
        </Panel>
        {selectedNode && !readOnly && (
          <Panel position="bottom-right" className="bg-white dark:bg-gray-800 p-2 rounded shadow">
            <button
              onClick={handleDeleteSelected}
              className="px-3 py-1 bg-red-500 text-white rounded hover:bg-red-600 text-sm"
            >
              Delete Node
            </button>
          </Panel>
        )}
      </ReactFlow>
    </div>
  );
}