ai-chatbot/components/flow-board/mind-annotation-board.ts

import {
  addSelectedElement,
  BoardTransforms,
  clearSelectedElement,
  PlaitBoard,
  type PlaitElement,
  PlaitHistoryBoard,
  PlaitPointerType,
  type Point,
  Transforms,
} from "@plait/core";
import {
  ArrowLineMarkerType,
  ArrowLineShape,
  BasicShapes,
  createArrowLineElement,
  createGeometryElementWithoutText,
} from "@plait/draw";
import { getRectangleByNode, MindElement } from "@plait/mind";
import {
  type AnnotatedElement,
  getBoundaryPoints,
  getRelationshipPoints,
  isMindAnnotation,
  type MindAnnotation,
  type TopicRectangle,
} from "@/lib/flow-board/mind-annotations";

function getTopics(board: PlaitBoard) {
  const topics = new Map<
    string,
    { element: MindElement; visible: MindElement }
  >();
  const visit = (elements: PlaitElement[], collapsedAncestor?: MindElement) => {
    for (const element of elements) {
      if (!MindElement.isMindElement(board, element)) continue;
      topics.set(element.id, {
        element,
        visible: collapsedAncestor || element,
      });
      visit(
        element.children || [],
        collapsedAncestor || (element.isCollapsed ? element : undefined)
      );
    }
  };
  visit(board.children);
  return topics;
}

function rectangle(element: MindElement): TopicRectangle | null {
  const node = MindElement.getNode(element);
  return node ? getRectangleByNode(node) : null;
}

function annotationPoints(
  board: PlaitBoard,
  annotation: MindAnnotation
): Point[] | null {
  const topics = getTopics(board);
  if (annotation.kind === "relationship") {
    const source = topics.get(annotation.sourceId);
    const target = topics.get(annotation.targetId);
    if (!source || !target) return null;
    const a = rectangle(source.visible);
    const b = rectangle(target.visible);
    return a && b ? getRelationshipPoints(a, b) : null;
  }
  const rectangles: TopicRectangle[] = [];
  const seen = new Set<string>();
  const collect = (element: MindElement) => {
    if (seen.has(element.id)) return;
    seen.add(element.id);
    const rect = rectangle(element);
    if (rect) rectangles.push(rect);
    if (!element.isCollapsed) element.children?.forEach(collect);
  };
  for (const id of annotation.topicIds) {
    const topic = topics.get(id);
    if (topic) collect(topic.visible);
  }
  return getBoundaryPoints(rectangles, annotation.padding);
}

export function insertMindAnnotation(
  board: PlaitBoard,
  annotation: MindAnnotation
) {
  const points = annotationPoints(board, annotation);
  if (!points) return null;
  const endpoints: [Point, Point] = [points[0], points[points.length - 1]];
  const label = { type: "paragraph", children: [{ text: "联系" }] };
  const shape =
    annotation.kind === "relationship"
      ? createArrowLineElement(
          ArrowLineShape.curve,
          endpoints,
          { marker: ArrowLineMarkerType.none },
          { marker: ArrowLineMarkerType.arrow },
          [{ position: 0.5, text: label }],
          { strokeColor: "#64748b", strokeWidth: 2 }
        )
      : createGeometryElementWithoutText(
          BasicShapes.roundRectangle,
          endpoints,
          {
            fill: "transparent",
            strokeColor: "#8b9bb8",
            strokeWidth: 2,
          }
        );
  const element: AnnotatedElement = {
    ...shape,
    points,
    mindAnnotation: annotation,
  };
  PlaitHistoryBoard.withNewBatch(board, () => {
    // Boundaries sit behind topics, so clicking within them still selects topics.
    Transforms.insertNode(board, element, [
      annotation.kind === "boundary" ? 0 : board.children.length,
    ]);
  });
  BoardTransforms.updatePointerType(board, PlaitPointerType.selection);
  clearSelectedElement(board);
  addSelectedElement(board, element);
  return element;
}

const installed = new WeakSet<PlaitBoard>();

/** Keep native drawable shapes linked to topic IDs, using the existing undo stack. */
export function withMindAnnotations(board: PlaitBoard) {
  if (installed.has(board)) return;
  installed.add(board);
  const { afterChange, undo, redo, apply } = board;
  board.apply = (operation) => {
    // Plait keeps a hover reference after deleting or collapsing a topic. Clear it
    // while the node's DOM still exists, otherwise the next pointerLeave throws.
    if (
      operation.type === "remove_node" ||
      (operation.type === "set_node" &&
        ("isCollapsed" in operation.newProperties ||
          "children" in operation.newProperties))
    ) {
      board.pointerLeave(new PointerEvent("pointerleave"));
    }
    apply(operation);
  };
  let pending = false;
  let replaying = true;
  const synchronize = () => {
    pending = false;
    const replay = replaying;
    replaying = false;
    const topics = getTopics(board);
    const changes = board.children
      .filter(isMindAnnotation)
      .flatMap((element) => {
        const annotation = element.mindAnnotation;
        const missing =
          annotation.kind === "relationship"
            ? !topics.has(annotation.sourceId) ||
              !topics.has(annotation.targetId)
            : !annotation.topicIds.some((id) => topics.has(id));
        const points = missing ? null : annotationPoints(board, annotation);
        if (
          !missing &&
          (!points || JSON.stringify(points) === JSON.stringify(element.points))
        )
          return [];
        return [{ element, points, missing }];
      });
    if (!changes.length) return;
    const update = () => {
      for (const { element, points, missing } of changes) {
        const index = board.children.findIndex(
          (child) => child.id === element.id
        );
        if (index < 0) continue;
        if (missing) Transforms.removeNode(board, [index]);
        else if (points) Transforms.setNode(board, { points }, [index]);
      }
    };
    if (replay) PlaitHistoryBoard.withoutSaving(board, update);
    else PlaitHistoryBoard.withMerging(board, update);
  };
  board.undo = () => {
    replaying = board.history.undos.length > 0;
    undo();
  };
  board.redo = () => {
    replaying = board.history.redos.length > 0;
    redo();
  };
  board.afterChange = () => {
    afterChange();
    if (
      !pending &&
      board.operations.some((op) => op.type !== "set_selection")
    ) {
      pending = true;
      queueMicrotask(synchronize);
    }
  };
  queueMicrotask(synchronize);
}