Socrates/apps/web/components/diagram-canvas/DiagramCanvas.tsx

// React Flow-backed diagram canvas (M3).
// - Custom block/actor/constraint/system nodes
// - Custom edges (association/composition/constraint)
// - Drag-create from the left palette
// - Click-edit via the right inspector
// - Drag handles between nodes to connect
//
// State is in-memory for M3; M5 swaps it for the bidirectional sync engine.

"use client";

import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import {
  ReactFlow,
  ReactFlowProvider,
  Background,
  BackgroundVariant,
  Controls,
  MarkerType,
  applyNodeChanges,
  applyEdgeChanges,
  addEdge,
  useReactFlow,
  type Node,
  type Edge,
  type NodeChange,
  type EdgeChange,
  type Connection,
  type NodeMouseHandler,
} from "@xyflow/react";
import { BlockNode, type BlockNodeData } from "./nodes/BlockNode";
import { SysmlEdge, type SysmlEdgeData } from "./edges/SysmlEdge";
import { Palette } from "./Palette";
import { NodeInspector } from "./NodeInspector";
import { fixtureToFlow } from "./fixtureToFlow";
import type { Density } from "../socrates/SocratesDock";
import type { FixtureData, AssociationKind, BlockKind } from "../../lib/fixtures/aristotle";
import type { ValidationIssue } from "../../lib/sysml/validate";

export type DiagramVariant = "softened" | "formal" | "graph";

interface DiagramCanvasProps {
  data: FixtureData;
  density?: Density;
  variant?: DiagramVariant;
  focusBlockId: string | null;
  onSelect?: (id: string | null) => void;
  issuesByElement?: Map<string, ValidationIssue[]>;
}

const nodeTypes = { sysmlBlock: BlockNode };
const edgeTypes = { sysml: SysmlEdge };

let nextNodeId = 1000;
let nextEdgeId = 1000;
function freshNodeId(): string { return `n${nextNodeId++}`; }
function freshEdgeId(): string { return `e${nextEdgeId++}`; }

function pickWorst(issues: ValidationIssue[]): "error" | "warning" | "soft" | undefined {
  if (issues.some(i => i.severity === "error")) return "error";
  if (issues.some(i => i.severity === "warning")) return "warning";
  if (issues.some(i => i.severity === "soft")) return "soft";
  return undefined;
}

export function DiagramCanvas(props: DiagramCanvasProps) {
  return (
    <ReactFlowProvider>
      <DiagramInner {...props} />
    </ReactFlowProvider>
  );
}

function DiagramInner({ data, focusBlockId, onSelect, issuesByElement }: DiagramCanvasProps) {
  const initial = useMemo(() => fixtureToFlow(data), [data]);
  const [nodes, setNodes] = useState<Node<BlockNodeData>[]>(initial.nodes);
  const [edges, setEdges] = useState<Edge<SysmlEdgeData>[]>(initial.edges);
  const wrapperRef = useRef<HTMLDivElement | null>(null);
  const { screenToFlowPosition } = useReactFlow();

  const onNodesChange = useCallback(
    (changes: NodeChange[]) => setNodes(ns => applyNodeChanges(changes, ns) as Node<BlockNodeData>[]),
    []
  );
  const onEdgesChange = useCallback(
    (changes: EdgeChange[]) => setEdges(es => applyEdgeChanges(changes, es) as Edge<SysmlEdgeData>[]),
    []
  );
  const onConnect = useCallback((connection: Connection) => {
    const newEdge: Edge<SysmlEdgeData> = {
      id: freshEdgeId(),
      source: connection.source,
      target: connection.target,
      type: "sysml",
      data: { label: "relates_to", kind: "association" satisfies AssociationKind },
      markerEnd: { type: MarkerType.Arrow, color: "var(--edge)", width: 18, height: 18 },
    };
    setEdges(es => addEdge(newEdge, es) as Edge<SysmlEdgeData>[]);
  }, []);

  // Highlight focusBlockId from sibling components (e.g. narrative chip hover).
  // Also surface validation severity onto each node so BlockNode can ring it.
  useEffect(() => {
    setNodes(ns =>
      ns.map(n => {
        const issues = issuesByElement?.get(n.id) ?? [];
        const severity = pickWorst(issues);
        return {
          ...n,
          selected: n.id === focusBlockId,
          data: { ...(n.data as BlockNodeData), issueSeverity: severity },
        };
      })
    );
  }, [focusBlockId, issuesByElement]);

  const onNodeClick: NodeMouseHandler = useCallback(
    (_event, node) => {
      onSelect?.(node.id);
    },
    [onSelect]
  );

  const onPaneClick = useCallback(() => {
    onSelect?.(null);
  }, [onSelect]);

  const selectedNode = nodes.find(n => n.id === focusBlockId);

  // Drag-and-drop from the palette
  const onDragOver = useCallback((event: React.DragEvent) => {
    event.preventDefault();
    event.dataTransfer.dropEffect = "move";
  }, []);

  const onDrop = useCallback(
    (event: React.DragEvent) => {
      event.preventDefault();
      const kindRaw = event.dataTransfer.getData("application/sysml-kind");
      if (!kindRaw) return;
      const kind = kindRaw as BlockKind | "system";
      const position = screenToFlowPosition({ x: event.clientX, y: event.clientY });
      const id = freshNodeId();
      const newNode: Node<BlockNodeData> = {
        id,
        type: "sysmlBlock",
        position,
        data: {
          label: kind === "system"
            ? "System"
            : kind === "actor"
            ? "Actor"
            : kind === "constraint"
            ? "Constraint"
            : "Block",
          kind,
          properties: kind === "constraint" ? [] : ["new_property"],
          ...(kind === "constraint" ? { expression: "{ }" } : {}),
        },
      };
      setNodes(ns => [...ns, newNode]);
      onSelect?.(id);
    },
    [onSelect, screenToFlowPosition]
  );

  function patchSelected(patch: Partial<BlockNodeData>) {
    if (!focusBlockId) return;
    setNodes(ns =>
      ns.map(n => (n.id === focusBlockId ? { ...n, data: { ...(n.data as BlockNodeData), ...patch } } : n))
    );
  }

  function deleteSelected() {
    if (!focusBlockId) return;
    setNodes(ns => ns.filter(n => n.id !== focusBlockId));
    setEdges(es => es.filter(e => e.source !== focusBlockId && e.target !== focusBlockId));
    onSelect?.(null);
  }

  return (
    <div ref={wrapperRef} className="diagram-flow" onDragOver={onDragOver} onDrop={onDrop}>
      <ReactFlow
        nodes={nodes}
        edges={edges}
        nodeTypes={nodeTypes}
        edgeTypes={edgeTypes}
        onNodesChange={onNodesChange}
        onEdgesChange={onEdgesChange}
        onConnect={onConnect}
        onNodeClick={onNodeClick}
        onPaneClick={onPaneClick}
        fitView
        fitViewOptions={{ padding: 0.15, includeHiddenNodes: false, maxZoom: 1.2 }}
        proOptions={{ hideAttribution: true }}
        defaultEdgeOptions={{
          type: "sysml",
          markerEnd: { type: MarkerType.Arrow, color: "var(--edge)", width: 18, height: 18 },
        }}
      >
        <Background variant={BackgroundVariant.Dots} gap={18} size={1} color="var(--grid-dot)" />
        <Controls showInteractive={false} position="bottom-right" />
      </ReactFlow>

      <Palette />

      {selectedNode && (
        <NodeInspector
          nodeId={selectedNode.id}
          data={selectedNode.data as BlockNodeData}
          onChange={patchSelected}
          onDelete={deleteSelected}
          onClose={() => onSelect?.(null)}
        />
      )}
    </div>
  );
}