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

// React Flow-backed diagram canvas (M5).
//
// State separation:
//   - SysMLModel  → ModelStore (canonical, semantic). The narrative reads here too.
//   - React Flow  → owns positions, selection, drag-in-progress (UI ephemeral).
//
// Sync rules:
//   - On model change from outside (chip rename, etc.), we apply targeted
//     `setNodes` updates (label, kind, properties) WITHOUT touching positions.
//     Newly-added blocks/constraints get added; removed ones get removed.
//   - On user actions in RF (connect, drop-from-palette, node delete, edge
//     delete, inspector edit), we emit ModelOps via useApply().
//
// This avoids the "re-derive nodes on every drag tick" trap that made
// dragging feel laggy in the first cut.

"use client";

import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import {
  ReactFlow,
  ReactFlowProvider,
  Background,
  BackgroundVariant,
  Controls,
  MarkerType,
  useNodesState,
  useEdgesState,
  useReactFlow,
  type Edge,
  type Node,
  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 { useModel, useApply } from "../../lib/sync/ModelStore";
import {
  addBlock as addBlockOp,
  addConstraint as addConstraintOp,
  removeBlock as removeBlockOp,
  removeConstraint as removeConstraintOp,
  removeAssociation as removeAssociationOp,
  updateBlock as updateBlockOp,
  updateConstraint as updateConstraintOp,
  addAssociation as addAssociationOp,
  addProperty as addPropertyOp,
  updateProperty as updatePropertyOp,
  removeProperty as removePropertyOp,
  newProperty,
  tempId,
  type ModelOp,
} from "../../lib/sync/ops";
import type { ValidationIssue } from "../../lib/sysml/validate";
import type { Density } from "../socrates/SocratesDock";
import type { FixtureData, BlockKind } from "../../lib/fixtures/aristotle";
import type { Block, Property, PropertyType, SysMLModel } from "../../lib/sysml/model";

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 };

const BOARD_W = 720;
const BOARD_H = 460;

function pickWorst(issues: ValidationIssue[] | undefined): "error" | "warning" | "soft" | undefined {
  if (!issues) return 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 model = useModel();
  const apply = useApply();
  const wrapperRef = useRef<HTMLDivElement | null>(null);
  const { screenToFlowPosition } = useReactFlow();

  // Initialize React Flow state from the model + fixture positions on first render.
  // Lazy `useState`-style init to avoid recomputing on every render.
  const [initialN] = useState(() => initialNodes(model, data));
  const [initialE] = useState(() => initialEdges(model));
  const [nodes, setNodes, onNodesChangeRaw] = useNodesState<Node<BlockNodeData>>(initialN);
  const [edges, setEdges, onEdgesChangeRaw] = useEdgesState<Edge<SysmlEdgeData>>(initialE);

  // Track the model's element ids so we can detect adds/removes between renders.
  const elementIdsRef = useRef<{ blocks: Set<string>; assocs: Set<string>; constraints: Set<string> }>(snapshotIds(model));

  // ─── Model → React Flow (targeted sync) ──────────────────────────────────
  useEffect(() => {
    const prev = elementIdsRef.current;
    const curr = snapshotIds(model);

    setNodes(currentNodes => {
      const byId = new Map(currentNodes.map(n => [n.id, n]));
      const next: Node<BlockNodeData>[] = [];

      // Existing model elements first — preserve position & selection from RF state
      for (const b of model.blocks) {
        const existing = byId.get(b.id);
        if (existing) {
          // Update data only if it changed (cheap structural compare)
          const propNames = b.properties.map(p => p.name);
          const dataChanged =
            existing.data?.label !== b.label ||
            existing.data?.kind !== b.kind ||
            !sameStringArray(existing.data?.properties ?? [], propNames);
          if (dataChanged) {
            next.push({ ...existing, data: { ...existing.data, label: b.label, kind: b.kind, properties: propNames } });
          } else {
            next.push(existing);
          }
        } else {
          // New block from outside (rare in M5; mostly self-originated drops)
          next.push(makeNodeForBlock(b));
        }
      }
      for (const c of model.constraints) {
        const existing = byId.get(c.id);
        if (existing) {
          const expr = c.expression || "{ }";
          const dataChanged =
            existing.data?.label !== c.label ||
            existing.data?.expression !== expr;
          if (dataChanged) {
            next.push({ ...existing, data: { ...existing.data, label: c.label, kind: "constraint", properties: [], expression: expr } });
          } else {
            next.push(existing);
          }
        } else {
          next.push(makeNodeForConstraint(c));
        }
      }
      return next;
    });

    setEdges(currentEdges => {
      const byId = new Map(currentEdges.map(e => [e.id, e]));
      const next: Edge<SysmlEdgeData>[] = [];

      for (const a of model.associations) {
        const existing = byId.get(a.id);
        if (existing && existing.data?.label === a.label && existing.data?.kind === a.kind) {
          next.push(existing);
        } else {
          next.push(makeEdgeForAssociation(a));
        }
      }
      // Synthesized constraint→block edges
      for (const c of model.constraints) {
        for (const target of c.appliesTo) {
          const id = `${c.id}__applies__${target}`;
          const existing = byId.get(id);
          next.push(existing ?? makeEdgeForApplies(c.id, target));
        }
      }
      return next;
    });

    elementIdsRef.current = curr;
    void prev;
  }, [model, setNodes, setEdges]);

  // ─── Selection sync — narrative chip hover / rail entry hover → here ────
  useEffect(() => {
    setNodes(curr => {
      let changed = false;
      const next = curr.map(n => {
        const shouldSelect = n.id === focusBlockId;
        if (n.selected === shouldSelect) return n;
        changed = true;
        return { ...n, selected: shouldSelect };
      });
      return changed ? next : curr;
    });
  }, [focusBlockId, setNodes]);

  // ─── Issue-severity sync — validation results → node ring color ─────────
  useEffect(() => {
    setNodes(curr => {
      let changed = false;
      const next = curr.map(n => {
        const next = pickWorst(issuesByElement?.get(n.id));
        if (n.data?.issueSeverity === next) return n;
        changed = true;
        return { ...n, data: { ...n.data, issueSeverity: next } };
      });
      return changed ? next : curr;
    });
  }, [issuesByElement, setNodes]);

  // ─── React Flow → ModelStore handlers ───────────────────────────────────

  const onNodesChange = useCallback((changes: NodeChange[]) => {
    // 1. Let React Flow update its own state (positions, drags, selections)
    onNodesChangeRaw(changes as NodeChange<Node<BlockNodeData>>[]);

    // 2. For removals, emit the corresponding remove ops
    const removalOps: ModelOp[] = [];
    for (const c of changes) {
      if (c.type !== "remove") continue;
      const isConstraint = !!model.constraints.find(x => x.id === c.id);
      if (isConstraint) {
        removalOps.push(removeConstraintOp(c.id));
      } else if (model.blocks.find(x => x.id === c.id)) {
        removalOps.push(removeBlockOp(c.id));
      }
    }
    if (removalOps.length > 0) apply(removalOps);
  }, [apply, model.blocks, model.constraints, onNodesChangeRaw]);

  const onEdgesChange = useCallback((changes: EdgeChange[]) => {
    // 1. Let React Flow update its own state
    onEdgesChangeRaw(changes as EdgeChange<Edge<SysmlEdgeData>>[]);

    // 2. For removals, emit ops — except for synthesized constraint-applies
    //    edges which decompose into a constraint update.
    const removalOps: ModelOp[] = [];
    for (const c of changes) {
      if (c.type !== "remove") continue;

      if (c.id.includes("__applies__")) {
        const [cid, tid] = c.id.split("__applies__");
        const target = model.constraints.find(x => x.id === cid);
        if (target) {
          removalOps.push(updateConstraintOp(cid!, { appliesTo: target.appliesTo.filter(x => x !== tid) }));
        }
      } else if (model.associations.find(a => a.id === c.id)) {
        removalOps.push(removeAssociationOp(c.id));
      }
    }
    if (removalOps.length > 0) apply(removalOps);
  }, [apply, model.associations, model.constraints, onEdgesChangeRaw]);

  const onConnect = useCallback((connection: Connection) => {
    if (!connection.source || !connection.target) return;
    const newId = tempId("a");
    apply([
      addAssociationOp({
        id: newId,
        fromBlockId: connection.source,
        toBlockId: connection.target,
        label: "relates_to",
        kind: "association",
      }, newId),
    ]);
  }, [apply]);

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

  const selectedBlock = useMemo(
    () => model.blocks.find(b => b.id === focusBlockId) ?? null,
    [model.blocks, focusBlockId]
  );
  const selectedConstraint = useMemo(
    () => model.constraints.find(c => c.id === focusBlockId) ?? null,
    [model.constraints, 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 position = screenToFlowPosition({ x: event.clientX, y: event.clientY });

      if (kindRaw === "constraint") {
        const cid = tempId("c");
        const result = apply([
          addConstraintOp({ id: cid, label: "Constraint", expression: "{ }", appliesTo: [] }, cid),
        ]);
        if (result.applied) {
          const final = result.idMapping[cid] ?? cid;
          // Position the new node at the drop coordinates
          setNodes(curr => curr.map(n => (n.id === final ? { ...n, position } : n)));
          onSelect?.(final);
        }
        return;
      }

      const kind = kindRaw as BlockKind | "system";
      const bid = tempId("b");
      const block: Block = {
        id: bid,
        label:
          kind === "system" ? "System" :
          kind === "actor" ? "Actor" :
          "Block",
        kind: kind === "system" ? "system" : kind,
        stereotypes: [kind],
        properties: [newProperty("new_property")],
      };
      const result = apply([addBlockOp(block, bid)]);
      if (result.applied) {
        const final = result.idMapping[bid] ?? bid;
        setNodes(curr => curr.map(n => (n.id === final ? { ...n, position } : n)));
        onSelect?.(final);
      }
    },
    [apply, onSelect, screenToFlowPosition, setNodes]
  );

  function patchSelected(patch: Partial<BlockNodeData>) {
    if (selectedBlock) {
      const ops: ModelOp[] = [];
      if (patch.label !== undefined && patch.label !== selectedBlock.label) {
        ops.push(updateBlockOp(selectedBlock.id, { label: patch.label }));
      }
      if (patch.kind !== undefined && patch.kind !== selectedBlock.kind) {
        ops.push(updateBlockOp(selectedBlock.id, { kind: patch.kind, stereotypes: [patch.kind] }));
      }
      if (patch.properties !== undefined) {
        ops.push(...diffProperties(selectedBlock, patch.properties));
      }
      if (ops.length > 0) apply(ops);
    } else if (selectedConstraint) {
      const ops: ModelOp[] = [];
      if (patch.label !== undefined && patch.label !== selectedConstraint.label) {
        ops.push(updateConstraintOp(selectedConstraint.id, { label: patch.label }));
      }
      if (patch.expression !== undefined && patch.expression !== selectedConstraint.expression) {
        ops.push(updateConstraintOp(selectedConstraint.id, { expression: patch.expression }));
      }
      if (ops.length > 0) apply(ops);
    }
  }

  function deleteSelected() {
    if (selectedBlock) {
      apply([removeBlockOp(selectedBlock.id)]);
      onSelect?.(null);
    } else if (selectedConstraint) {
      apply([removeConstraintOp(selectedConstraint.id)]);
      onSelect?.(null);
    }
  }

  const inspectorData: BlockNodeData | null = selectedBlock
    ? {
        label: selectedBlock.label,
        kind: selectedBlock.kind,
        properties: selectedBlock.properties.map(p => p.name),
      }
    : selectedConstraint
    ? {
        label: selectedConstraint.label,
        kind: "constraint",
        properties: [],
        expression: selectedConstraint.expression,
      }
    : 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}
        deleteKeyCode={["Backspace", "Delete"]}
        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 />

      {inspectorData && focusBlockId && (
        <NodeInspector
          nodeId={focusBlockId}
          data={inspectorData}
          onChange={patchSelected}
          onDelete={deleteSelected}
          onClose={() => onSelect?.(null)}
        />
      )}
    </div>
  );
}

// ─── Initial state builders ──────────────────────────────────────────────

function initialNodes(model: SysMLModel, data?: FixtureData): Node<BlockNodeData>[] {
  // Bootstrap positions from the fixture for the original Aristotle blocks.
  const fixturePositions: Record<string, { x: number; y: number }> = {};
  if (data) {
    for (const b of data.blocks) {
      fixturePositions[b.id] = { x: b.x * BOARD_W, y: b.y * BOARD_H };
    }
  }
  const nodes: Node<BlockNodeData>[] = [];
  let autoCol = 0;
  let autoRow = 0;
  function autoPosition() {
    const pos = { x: 60 + autoCol * 220, y: 60 + autoRow * 160 };
    autoCol++;
    if (autoCol >= 4) { autoCol = 0; autoRow++; }
    return pos;
  }

  for (const b of model.blocks) {
    nodes.push({
      ...makeNodeForBlock(b),
      position: fixturePositions[b.id] ?? autoPosition(),
    });
  }
  for (const c of model.constraints) {
    nodes.push({
      ...makeNodeForConstraint(c),
      position: fixturePositions[c.id] ?? autoPosition(),
    });
  }
  return nodes;
}

function initialEdges(model: SysMLModel): Edge<SysmlEdgeData>[] {
  const edges: Edge<SysmlEdgeData>[] = [];
  for (const a of model.associations) edges.push(makeEdgeForAssociation(a));
  for (const c of model.constraints) {
    for (const target of c.appliesTo) {
      edges.push(makeEdgeForApplies(c.id, target));
    }
  }
  return edges;
}

function makeNodeForBlock(b: Block): Node<BlockNodeData> {
  return {
    id: b.id,
    type: "sysmlBlock",
    position: { x: 0, y: 0 },
    data: {
      label: b.label,
      kind: b.kind,
      properties: b.properties.map(p => p.name),
    },
  };
}

function makeNodeForConstraint(c: import("../../lib/sysml/model").Constraint): Node<BlockNodeData> {
  return {
    id: c.id,
    type: "sysmlBlock",
    position: { x: 0, y: 0 },
    data: {
      label: c.label,
      kind: "constraint",
      properties: [],
      expression: c.expression || "{ }",
    },
  };
}

function makeEdgeForAssociation(a: import("../../lib/sysml/model").Association): Edge<SysmlEdgeData> {
  return {
    id: a.id,
    source: a.fromBlockId,
    target: a.toBlockId,
    type: "sysml",
    data: { label: a.label, kind: a.kind },
    markerEnd: a.kind === "composition"
      ? { type: MarkerType.ArrowClosed, color: "var(--edge)", width: 18, height: 18 }
      : a.kind === "constraintApplies"
      ? undefined
      : { type: MarkerType.Arrow, color: "var(--edge)", width: 18, height: 18 },
  };
}

function makeEdgeForApplies(constraintId: string, targetBlockId: string): Edge<SysmlEdgeData> {
  return {
    id: `${constraintId}__applies__${targetBlockId}`,
    source: constraintId,
    target: targetBlockId,
    type: "sysml",
    data: { label: "applies_to", kind: "constraintApplies" },
  };
}

function snapshotIds(model: SysMLModel): { blocks: Set<string>; assocs: Set<string>; constraints: Set<string> } {
  return {
    blocks: new Set(model.blocks.map(b => b.id)),
    assocs: new Set(model.associations.map(a => a.id)),
    constraints: new Set(model.constraints.map(c => c.id)),
  };
}

function sameStringArray(a: string[], b: string[]): boolean {
  if (a.length !== b.length) return false;
  for (let i = 0; i < a.length; i++) if (a[i] !== b[i]) return false;
  return true;
}

/** Diff property-name lists into add/update/remove ops. */
function diffProperties(block: Block, nextNames: string[]): ModelOp[] {
  const ops: ModelOp[] = [];
  const cur = block.properties;
  for (let i = 0; i < cur.length; i++) {
    const c = cur[i]!;
    const n = nextNames[i];
    if (n === undefined) {
      ops.push(removePropertyOp(block.id, c.id));
    } else if (n !== c.name) {
      ops.push(updatePropertyOp(block.id, c.id, { name: n }));
    }
  }
  for (let i = cur.length; i < nextNames.length; i++) {
    const n = nextNames[i]!;
    const pid = tempId("p");
    const p: Property = { id: pid, name: n, type: { kind: "string" } as PropertyType, multiplicity: "0..1" };
    ops.push(addPropertyOp(block.id, p, pid));
  }
  return ops;
}