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

/**
 * Build a Craft.js `NodeTree` from a `SerializedTreeNode` using
 * `query.parseFreshNode`. This is the correct way to construct a tree for
 * `actions.addNodeTree()` when inserting/replacing sections, individual
 * nodes, or (via `templates/apply-template.ts`) a template's header/footer/
 * page content -- anything whose input isn't already a live Craft.js state.
 *
 * Originally lived in `apply-ai-response.ts` (the AI-response path pioneered
 * it); moved here alongside `sanitizeAiTree`/`flattenTreeForCraft` once
 * `templates/apply-template.ts` needed the exact same tree-materialization
 * logic — building a node tree by hand via `React.createElement` +
 * `query.parseReactElement` does NOT correctly route content through a
 * `Section`/`BackgroundSection`/`FormContainer`'s internal SHELL_INNER
 * linked canvas or `ColumnLayout`'s linked columns (see the `SHELL_INNER`
 * comment above): Craft.js would auto-materialize that linked node at
 * render time with a botched `type` field and crash `toNodeTree` with
 * "component type (undefined) does not exist in the resolver". Going through
 * `sanitizeAiTree` + `flattenTreeForCraft` (as this function does) avoids
 * that entirely.
 */
export function buildNodeTree(query: any, tree: SerializedTreeNode): NodeTree {
  // Validate + sanitize at the boundary before touching Craft.js at all:
  // unknown resolvedNames are dropped (never reach parseFreshNode/addNodeTree,
  // which would throw), and ids are repaired ('ROOT'/empty/colliding →
  // regenerated) against the live document's existing node ids so a new
  // caller-supplied id can never clobber or duplicate one already on the canvas.
  const existingIds = new Set<string>(
    typeof query?.getNodes === 'function' ? Object.keys(query.getNodes()) : [],
  );
  const sanitized = sanitizeAiTree(tree, existingIds);
  if (!sanitized) {
    throw new Error('craft-tree: tree root has an unknown/invalid component type; nothing to build');
  }

  // sanitizeAiTree() has already guaranteed every node has a valid, unique,
  // non-'ROOT' node_id and a registered resolvedName — flattenTreeForCraft
  // just needs to do the structural walk (linkedNodes/SHELL_INNER/style
  // normalization) common to every tree-flattening caller in this codebase.
  const flat = flattenTreeForCraft(sanitized);

  const craftNodes: Record<string, any> = {};
  for (const [id, flatNode] of Object.entries(flat.nodes)) {
    // Live Craft.js nodes need `data.type` to be the ACTUAL component
    // reference (the function/class from `componentResolver`) -- not the
    // `{ resolvedName }` wrapper object `flattenTreeForCraft` produces (that
    // wrapper is the *serialized* form used by `actions.deserialize()`/
    // `query.serialize()`, a different Craft.js entry point), and not a
    // plain resolvedName string either.
    //
    // A plain string passes `parseFreshNode`'s own validation (it special-
    // cases strings) AND `actions.addNodeTree`'s, so that alone looks fine --
    // but it's wrong for RENDERING: Craft.js's node renderer does
    // `React.createElement(data.type, props)` directly, so a string type
    // makes React treat e.g. "Section" as a literal (invalid, lowercase-
    // expected) HTML tag name instead of resolving it to the real `Section`
    // component. Symptom: "is using incorrect casing" console warnings and
    // the component's own render logic (and thus its content) never runs --
    // reproduced live by loading a template: the header/footer rendered fine
    // but the page canvas went blank. The `{ resolvedName }` wrapper object
    // fails even earlier, at `parseFreshNode`/`addNodeTree`'s own validation
    // ("component type (undefined) does not exist in the resolver"), because
    // their resolution only accepts a string or a component reference.
    //
    // None of this was caught by the pre-existing mocked test for this
    // function (in apply-ai-response.test.ts) because its fake
    // `parseFreshNode` just echoes the input back unchanged rather than
    // exercising Craft.js's real validation/rendering.
    const resolvedComponent = (componentResolver as Record<string, unknown>)[flatNode.type.resolvedName];

    if (flat.syntheticIds.has(id)) {
      // The SHELL_INNER wrapper Container is synthesized by the flattener,
      // not present in the input tree — build it directly rather than
      // through query.parseFreshNode(), which would merge in Container's
      // default `craft.props` (e.g. default padding) and change its look.
      //
      // `actions.addNodeTree`'s own validation for a non-string `data.type`
      // checks `resolver[data.name]` (NOT a reverse lookup from `data.type`
      // like `parseFreshNode` does) -- so `data.name` must be set to the
      // resolvedName string here too, or this synthetic node fails that
      // check with "component type (Container) does not exist in the
      // resolver" the moment `addNodeTree` runs, even though `resolvedComponent`
      // itself is perfectly valid.
      craftNodes[id] = {
        id,
        data: { ...flatNode, type: resolvedComponent, name: flatNode.type.resolvedName },
        events: { selected: false, hovered: false, dragged: false },
        rules: { canDrag: () => true, canMoveIn: () => true, canMoveOut: () => true, canDrop: () => true },
      };
    } else {
      craftNodes[id] = (query.parseFreshNode({
        id,
        data: { ...flatNode, type: resolvedComponent },
      }) as any).toNode();
    }
  }

  return { rootNodeId: flat.rootNodeId, nodes: craftNodes };
}

/**
 * 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 only the mutable sub-objects (props, custom) so they're
    // 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.
    //
    // Deliberately NOT `structuredClone(oldNode.data)` as a whole: for a
    // LIVE Craft.js node (as opposed to a plain serialized one), `data.type`
    // is the actual component function/class reference (see the long
    // comment on `buildNodeTree` above) -- `structuredClone` cannot clone a
    // function and throws `DataCloneError`, which silently broke EVERY
    // duplicate/paste in the real app (masked in unit tests that mock
    // `@craftjs/core` with plain-data fake nodes, so `toNodeTree()` never
    // actually returns a function there). `type` is a stable reference that
    // both the original and the duplicate should point at unchanged, so it
    // never needed cloning in the first place.
    const clonedData = {
      ...oldNode.data,
      props: structuredClone(oldNode.data.props),
      custom: structuredClone(oldNode.data.custom),
    };

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