site-builder/craft/src/utils/craft-tree.ts

import type { NodeTree, Node, NodeId } from '@craftjs/core';
import { getRandomId } from '@craftjs/utils';
import type { SerializedTreeNode } from '../types/sitesmith';
import { componentResolver } from '../components/resolver';

/** Only Container is a "real" Craft.js canvas in serialized state. Layout
 * shells (Section/HeroSimple/ColumnLayout/etc) use <Element canvas> linkedNodes
 * internally — their own node must serialize with isCanvas:false or
 * toNodeTree's walker hits an Invariant because the shell claims to be a
 * canvas but its render ignores `data.nodes`. */
export const CANVAS_TYPES = new Set(['Container']);

/** Shells that wrap their content in a single <Element id="<key>" is={Container}>.
 * When a tree puts content directly under one of these, the children end up
 * orphaned (the shell ignores data.nodes — it renders via the linkedNode) and
 * Craft.js auto-creates the linkedNode at render time with a botched type
 * field, which then crashes toNodeTree. Pre-create the linkedNode ourselves
 * to keep the state shape Craft.js expects. */
export const SHELL_INNER: Record<string, string> = {
  Section: 'section-inner',
  BackgroundSection: 'bg-section-inner',
  FormContainer: 'form-inner',
};

/** True if `name` is a registered component in the resolver (i.e. safe to
 * hand to `actions.addNodeTree`/`actions.deserialize` without it throwing
 * "does not exist in resolver" at render time). */
function isKnownResolvedName(name: unknown): name is string {
  return typeof name === 'string' && Object.prototype.hasOwnProperty.call(componentResolver, name);
}

/**
 * Validate + sanitize a tree (AI-supplied, or otherwise untrusted) before it
 * is handed to `flattenTreeForCraft`/`query.parseFreshNode`/`actions.deserialize`.
 * Fails SOFT — never throws:
 *
 *  - **resolvedName allowlist**: any node whose `type.resolvedName` is not a
 *    key of `componentResolver` is dropped (it and its subtree), with a
 *    `console.warn`. If the tree's ROOT is itself invalid, this returns
 *    `null` and the caller must skip the whole op/tree (or substitute a safe
 *    fallback state).
 *  - **id policy — REGENERATE (not skip)**: a node id that is `'ROOT'`,
 *    empty/non-string, or collides with an id already in `usedIds` is
 *    replaced with a fresh `ai-auto-N` id (and a warning is logged). The
 *    node itself is kept — only structurally invalid *types* are dropped;
 *    a merely-bad *id* is repaired so we never lose otherwise-valid content
 *    over an id collision.
 *
 * `usedIds` should be seeded with any ids that must not collide (e.g. a live
 * document's existing node ids); it is also used internally to keep ids
 * unique within the tree being sanitized.
 */
export function sanitizeAiTree(
  tree: SerializedTreeNode,
  usedIds: Set<string> = new Set(),
  idCounter: { n: number } = { n: 0 },
): SerializedTreeNode | null {
  const resolvedName = tree?.type?.resolvedName;
  if (!isKnownResolvedName(resolvedName)) {
    console.warn(`sitesmith: dropping node with unknown component "${String(resolvedName)}"`);
    return null;
  }

  const rawId = tree.props?.node_id;
  let id = typeof rawId === 'string' ? rawId.trim() : '';
  if (!id) {
    id = `ai-auto-${idCounter.n++}`;
  } else if (id === 'ROOT') {
    console.warn('sitesmith: rejecting AI-supplied node id "ROOT" (reserved); regenerating');
    id = `ai-auto-${idCounter.n++}`;
  } else if (usedIds.has(id)) {
    console.warn(`sitesmith: AI-supplied node id "${id}" collides with an existing id; regenerating`);
    id = `ai-auto-${idCounter.n++}`;
  }
  // Belt-and-suspenders: keep drawing fresh ids in the (pathological) case
  // the freshly generated id itself collides.
  while (usedIds.has(id)) id = `ai-auto-${idCounter.n++}`;
  usedIds.add(id);

  const children: SerializedTreeNode[] = [];
  for (const child of tree.nodes ?? []) {
    const sanitizedChild = sanitizeAiTree(child, usedIds, idCounter);
    if (sanitizedChild) children.push(sanitizedChild);
  }

  return {
    type: tree.type,
    props: { ...tree.props, node_id: id },
    nodes: children,
  };
}

/** A single flattened Craft.js serialized-node entry, as produced by
 * `flattenTreeForCraft`. Shape matches `@craftjs/core`'s `SerializedNode`. */
export interface FlatCraftNode {
  type: { resolvedName: string };
  isCanvas: boolean;
  props: Record<string, unknown>;
  displayName: string | undefined;
  custom: Record<string, unknown>;
  hidden: boolean;
  parent: string | null;
  nodes: string[];
  linkedNodes: Record<string, string>;
}

export interface FlattenTreeResult {
  rootNodeId: string;
  nodes: Record<string, FlatCraftNode>;
  /** ids of nodes synthesized by the flattener itself (currently: the
   *  auto-created SHELL_INNER wrapper for Section/BackgroundSection/
   *  FormContainer) rather than present in the input tree. Callers that
   *  materialize real Craft.js nodes (e.g. via `query.parseFreshNode`) need
   *  to know which ids these are so they can be built directly instead of
   *  run through the resolver (which would merge in the component's default
   *  `craft.props`, changing the synthetic wrapper's props). */
  syntheticIds: Set<string>;
}

/**
 * Flatten a `SerializedTreeNode` tree into a flat map of Craft.js
 * serialized-node entries, applying the transforms every caller in this
 * codebase needs:
 *
 *  - `style: []` -> `{}` (the AI sometimes emits an empty array instead of
 *    an empty object for a CSSProperties slot; React/Craft.js choke on that).
 *  - `ColumnLayout` children are moved from `nodes` into fixed `col-0,
 *    col-1, …` `linkedNodes` (ColumnLayout renders via linkedNodes, not
 *    direct children — direct children would be silently orphaned).
 *  - `Section`/`BackgroundSection`/`FormContainer` children are wrapped in a
 *    single pre-created `Container` linkedNode (see `SHELL_INNER`) so
 *    Craft.js never has to auto-materialize one at render time.
 *
 * This function assumes the input tree is already valid (every node's
 * `type.resolvedName` is a registered component, and `props.node_id` is a
 * unique, non-'ROOT' string where present) — callers should run untrusted
 * trees through `sanitizeAiTree` first. Ids are read from `props.node_id`,
 * falling back to an auto-generated id if absent.
 */
export function flattenTreeForCraft(tree: SerializedTreeNode): FlattenTreeResult {
  let counter = 0;
  const nodes: Record<string, FlatCraftNode> = {};
  const syntheticIds = new Set<string>();

  const walk = (node: SerializedTreeNode, parent: string | null): string => {
    const id = (node.props?.node_id as string | undefined) || `ai-auto-${counter++}`;
    const typeName = node.type?.resolvedName;
    const rawProps = node.props ?? {};
    const props: Record<string, unknown> = { ...rawProps };
    if (Array.isArray(props.style)) props.style = {};

    const childIds: string[] = [];
    nodes[id] = {
      type: node.type,
      // isCanvas must match the component's craft.rules — only layout
      // wrappers accept children. Setting it true on leaf components
      // (Heading, TextBlock, ButtonLink, etc) makes Craft.js render them
      // as empty drop-canvas wrappers and the actual content disappears.
      isCanvas: typeName ? CANVAS_TYPES.has(typeName) : false,
      props,
      displayName: typeName,
      custom: {},
      hidden: false,
      parent,
      nodes: childIds,
      linkedNodes: {},
    };
    for (const child of node.nodes ?? []) {
      childIds.push(walk(child, id));
    }

    // ColumnLayout uses Craft.js linkedNodes with fixed ids (col-0, col-1, ...).
    // Direct children would be orphaned by ColumnLayout's render (it creates
    // fresh column Elements and ignores data.nodes) and any subsequent
    // toNodeTree walk would hit an Invariant. Move them into linkedNodes so
    // they render in the columns actually shown to the user.
    if (typeName === 'ColumnLayout' && childIds.length > 0) {
      const linked: Record<string, string> = {};
      childIds.forEach((cid, i) => {
        linked[`col-${i}`] = cid;
        if (nodes[cid]) nodes[cid].isCanvas = true; // columns are canvases
      });
      nodes[id].nodes = [];
      nodes[id].linkedNodes = linked;
      const curProps = nodes[id].props;
      if (!curProps.columns || curProps.columns !== childIds.length) {
        curProps.columns = childIds.length;
      }
    }

    // Section/BackgroundSection/FormContainer each render a single
    // <Element id="<key>" is={Container} canvas> ... </Element>. If content
    // was nested as direct children, Craft.js would auto-create the
    // linkedNode on first render — storing its type as the Container
    // component class rather than {resolvedName:'Container'}, which then
    // crashes toNodeTree with "type (undefined) does not exist in resolver".
    // Pre-create the linkedNode ourselves with the correct serialized type
    // so Craft.js never has to materialize it.
    const innerKey = SHELL_INNER[typeName ?? ''];
    if (innerKey && childIds.length > 0) {
      const innerId = `${id}__${innerKey}`;
      nodes[innerId] = {
        type: { resolvedName: 'Container' },
        isCanvas: true,
        props: { tag: 'div' },
        displayName: 'Container',
        custom: {},
        hidden: false,
        parent: id,
        nodes: [...childIds],
        linkedNodes: {},
      };
      syntheticIds.add(innerId);
      for (const cid of childIds) {
        if (nodes[cid]) nodes[cid].parent = innerId;
      }
      nodes[id].nodes = [];
      nodes[id].linkedNodes = { [innerKey]: innerId };
    }

    return id;
  };

  const rootNodeId = walk(tree, null);
  return { rootNodeId, nodes, syntheticIds };
}

/**
 * Returns a deep-cloned NodeTree where every node id in the tree
 * (rootNodeId, the `nodes` map keys, each `node.id`, each internal
 * `node.data.parent`, every `node.data.nodes` entry, and every
 * `node.data.linkedNodes` value) has been remapped to a fresh,
 * collision-free id generated by Craft.js's own id generator.
 *
 * The subtree root's `data.parent` is deliberately left untouched --
 * `actions.addNodeTree(tree, parentId)` sets that itself when the tree is
 * attached to its new parent.
 *
 * The input tree is not mutated.
 */
export function regenerateTreeIds(tree: NodeTree): NodeTree {
  const oldIds = Object.keys(tree.nodes);

  // Build old -> new id map first so all cross-references can be rewritten
  // consistently regardless of iteration order.
  const idMap = new Map<NodeId, NodeId>();
  for (const oldId of oldIds) {
    idMap.set(oldId, getRandomId());
  }

  const remapId = (id: NodeId): NodeId => idMap.get(id) ?? id;

  const newNodes: Record<NodeId, Node> = {};
  for (const oldId of oldIds) {
    const oldNode = tree.nodes[oldId];
    const newId = idMap.get(oldId)!;

    const newParent =
      oldId === tree.rootNodeId
        ? oldNode.data.parent // subtree root's parent is set by the caller
        : oldNode.data.parent !== null && oldNode.data.parent !== undefined
          ? remapId(oldNode.data.parent)
          : oldNode.data.parent;

    const newLinkedNodes: Record<string, NodeId> = {};
    for (const [slot, linkedId] of Object.entries(oldNode.data.linkedNodes || {})) {
      newLinkedNodes[slot] = remapId(linkedId);
    }

    // Deep-clone data so mutable sub-objects (props, custom, etc.) are never
    // shared by reference between the original node and its regenerated
    // copy. Craft.js's setProp mutates data.props in place, so a shallow
    // copy here would let edits to the duplicate silently corrupt the
    // original.
    const clonedData = structuredClone(oldNode.data);

    const newNode: Node = {
      ...oldNode,
      id: newId,
      data: {
        ...clonedData,
        parent: newParent,
        nodes: (oldNode.data.nodes || []).map(remapId),
        linkedNodes: newLinkedNodes,
      },
    };

    newNodes[newId] = newNode;
  }

  return {
    rootNodeId: remapId(tree.rootNodeId),
    nodes: newNodes,
  };
}