site-builder/craft/src/utils/apply-ai-response.ts

import { useEditor } from '@craftjs/core';
import type { NodeTree } from '@craftjs/core';
import { usePages } from '../state/PageContext';
import { SitesmithResponse, 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`. */
const CANVAS_TYPES = new Set(['Container']);

const SHELL_INNER: Record<string, string> = {
  Section: 'section-inner',
  BackgroundSection: 'bg-section-inner',
  FormContainer: 'form-inner',
};

/**
 * update_props may never touch these keys, regardless of what the AI sends.
 * `node_id` is the stable identifier other ops (and findNodeIdByAiNodeId)
 * rely on to re-target a node in later turns of the conversation — letting
 * the AI clobber it would silently orphan future patches.
 */
const PROTECTED_UPDATE_PROPS_KEYS = new Set<string>(['node_id']);

/** True if `name` is a registered component in the resolver (i.e. safe to
 * hand to `actions.addNodeTree` 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 an AI-supplied tree before it is handed to
 * `buildNodeTree`/`query.parseFreshNode`. 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.
 *  - **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 AI
 *    content over an id collision.
 *
 * `usedIds` should be seeded with the live document's existing node ids
 * (see `buildNodeTree`) so AI-supplied ids can never collide with content
 * already on the canvas; 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,
  };
}

/**
 * Flatten a SerializedTreeNode tree into a Craft.js node map ready for
 * `actions.deserialize()`.
 *
 * Returns `{ rootNodeId, nodes }` where `nodes` is a flat map keyed by node id.
 * The root entry is also aliased under 'ROOT' so Craft.js can find it when
 * calling `actions.deserialize(JSON.stringify(nodes))`.
 */
export function serializeTreeForCraft(tree: SerializedTreeNode): { rootNodeId: string; nodes: Record<string, unknown> } {
  const idCounter = { n: 0 };
  const nodes: Record<string, any> = {};

  const walk = (node: SerializedTreeNode, parent: string | null): string => {
    const id = (node.props.node_id as string | undefined) || `ai-auto-${idCounter.n++}`;
    nodes[id] = {
      type: node.type,
      props: node.props,
      displayName: node.type.resolvedName,
      isCanvas: CANVAS_TYPES.has(node.type.resolvedName),
      parent,
      nodes: [] as string[],
      hidden: false,
      custom: {},
      linkedNodes: {},
    };
    for (const child of node.nodes ?? []) {
      const childId = walk(child, id);
      nodes[id].nodes.push(childId);
    }
    return id;
  };

  const rootId = walk(tree, null);

  // Craft.js frame expects a 'ROOT' key; alias it if the AI gave a different id
  if (rootId !== 'ROOT') {
    nodes['ROOT'] = { ...nodes[rootId], parent: null };
    // Fix children's parent reference to 'ROOT'
    for (const childId of nodes['ROOT'].nodes as string[]) {
      if (nodes[childId]) nodes[childId].parent = 'ROOT';
    }
  }

  return { rootNodeId: rootId, nodes };
}

/**
 * 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 or individual nodes.
 */
function buildNodeTree(query: any, tree: SerializedTreeNode): NodeTree {
  // Validate + sanitize at the AI 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
  // AI-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('sitesmith: AI tree root has an unknown/invalid component type; nothing to build');
  }

  const craftNodes: Record<string, any> = {};

  const walk = (node: SerializedTreeNode, parent: string | null): string => {
    // sanitizeAiTree() has already guaranteed every node here has a valid,
    // unique, non-'ROOT' node_id — no fallback/collision handling needed.
    const id = node.props.node_id as string;
    const craftNode = (query.parseFreshNode({
      id,
      data: {
        type: node.type,
        props: node.props,
        displayName: node.type.resolvedName,
        isCanvas: CANVAS_TYPES.has(node.type.resolvedName),
        parent,
        nodes: [],
        linkedNodes: {},
        hidden: false,
        custom: {},
      },
    }) as any).toNode() as any;
    craftNodes[id] = craftNode;
    for (const child of node.nodes ?? []) {
      const childId = walk(child, id);
      craftNodes[id].data.nodes.push(childId);
    }
    // ColumnLayout uses linkedNodes (col-0, col-1, ...) — not direct children.
    if (node.type.resolvedName === 'ColumnLayout' && craftNodes[id].data.nodes.length > 0) {
      const linked: Record<string, string> = {};
      craftNodes[id].data.nodes.forEach((cid: string, i: number) => {
        linked[`col-${i}`] = cid;
        if (craftNodes[cid]) craftNodes[cid].data.isCanvas = true;
      });
      craftNodes[id].data.nodes = [];
      craftNodes[id].data.linkedNodes = linked;
      const colCount = Object.keys(linked).length;
      if (!craftNodes[id].data.props.columns || craftNodes[id].data.props.columns !== colCount) {
        craftNodes[id].data.props.columns = colCount;
      }
    }
    // Section/BackgroundSection/FormContainer wrap their content in a single
    // <Element id="<key>" is={Container} canvas>. Pre-create the linkedNode
    // so Craft.js doesn't auto-create one with a malformed type field.
    const innerKey = SHELL_INNER[node.type.resolvedName];
    if (innerKey && craftNodes[id].data.nodes.length > 0) {
      const innerId = `${id}__${innerKey}`;
      const childIds: string[] = [...craftNodes[id].data.nodes];
      craftNodes[innerId] = {
        id: innerId,
        data: {
          type: { resolvedName: 'Container' },
          props: { tag: 'div' },
          displayName: 'Container',
          isCanvas: true,
          parent: id,
          nodes: childIds,
          linkedNodes: {},
          hidden: false,
          custom: {},
        },
        events: { selected: false, hovered: false, dragged: false },
        rules: { canDrag: () => true, canMoveIn: () => true, canMoveOut: () => true, canDrop: () => true },
      };
      for (const cid of childIds) {
        if (craftNodes[cid]) craftNodes[cid].data.parent = innerId;
      }
      craftNodes[id].data.nodes = [];
      craftNodes[id].data.linkedNodes = { [innerKey]: innerId };
    }
    return id;
  };

  const rootId = walk(sanitized, null);
  return { rootNodeId: rootId, nodes: craftNodes };
}

/**
 * Find the Craft.js node id that corresponds to an AI node_id value.
 * Checks `data.props.node_id` first, then falls back to raw id equality.
 */
export function findNodeIdByAiNodeId(query: any, aiNodeId: string): string | null {
  const all = query.getNodes() as Record<string, any>;
  for (const [id, n] of Object.entries(all)) {
    if (n.data?.props?.node_id === aiNodeId) return id;
    if (id === aiNodeId) return id;
  }
  return null;
}

/** Exported for unit tests */
export const __test = { findNodeIdByAiNodeId, buildNodeTree, applyPatch };

/**
 * React hook that returns an `apply` function.
 * Call `apply(response)` after a successful Sitesmith API call to materialize
 * the AI's instructions into the editor.
 */
export function useApplyAiResponse() {
  const { actions, query } = useEditor();
  const pages = usePages();

  /**
   * @param targetNodeId If set and the AI returned a section-scoped replace
   *   instead of a patch, treat the first returned tree as a replacement for
   *   this node (the user said "edit this block" — they don't want a new
   *   section appended at the bottom).
   */
  return async function apply(
    resp: SitesmithResponse,
    targetNodeId?: string,
  ): Promise<{ ok: boolean; message?: string }> {
    // 'ask' type = AI wants clarification, nothing to apply
    if (resp.type === 'ask') return { ok: true };

    if (resp.type === 'replace') {
      if (resp.scope === 'site') {
        pages.replaceAllPages(resp.pages.map((p) => ({ name: p.name, tree: p.tree })));
        if (resp.header) pages.setHeader(resp.header.tree);
        if (resp.footer) pages.setFooter(resp.footer.tree);
        return { ok: true };
      }

      if (resp.scope === 'page') {
        pages.replaceCurrentPage(resp.pages[0]);
        return { ok: true };
      }

      if (resp.scope === 'section') {
        // When targeted at a specific node, the AI's tree replaces that node
        // in place (vs appending a fresh section at the end of ROOT).
        if (targetNodeId && resp.pages.length > 0) {
          try {
            const nodeTree = buildNodeTree(query, resp.pages[0].tree);
            const parent: string = query.node(targetNodeId).get().data.parent ?? 'ROOT';
            const siblings: string[] = query.node(parent).childNodes();
            const index = siblings.indexOf(targetNodeId);
            actions.delete(targetNodeId);
            actions.addNodeTree(nodeTree, parent, index);
            return { ok: true };
          } catch (e) {
            console.warn('sitesmith: targeted section replace failed, falling back to append', e);
          }
        }
        // Insert each provided tree as a new node tree appended to ROOT
        for (const p of resp.pages) {
          try {
            const nodeTree = buildNodeTree(query, p.tree);
            actions.addNodeTree(nodeTree, 'ROOT');
          } catch (e) {
            console.warn('sitesmith: failed to add section tree', e);
          }
        }
        return { ok: true };
      }
    }

    if (resp.type === 'patch') {
      return applyPatch(actions, query, resp.ops);
    }

    return { ok: false, message: 'Unknown response type' };
  };
}

function applyPatch(
  actions: any,
  query: any,
  ops: any[],
): { ok: boolean; message?: string } {
  for (const op of ops) {
    const id = findNodeIdByAiNodeId(query, op.node_id);
    if (!id) {
      console.warn('sitesmith patch: node_id not found, skipping op:', op.node_id, op.op);
      continue;
    }

    switch (op.op) {
      case 'update_props': {
        const patchProps = op.props;
        if (!patchProps || typeof patchProps !== 'object' || Array.isArray(patchProps)) {
          console.warn('sitesmith patch: update_props with invalid props payload, skipping', op.node_id);
          break;
        }
        actions.setProp(id, (p: any) => {
          for (const [key, value] of Object.entries(patchProps)) {
            if (PROTECTED_UPDATE_PROPS_KEYS.has(key)) {
              console.warn(`sitesmith patch: refusing to overwrite protected prop "${key}"`);
              continue;
            }
            // Merge `style` shallowly instead of blindly replacing it — the
            // AI usually only means to change one or two style keys, and a
            // blind Object.assign would clobber every other style the user
            // (or a previous AI turn) had already set.
            if (key === 'style' && value && typeof value === 'object' && !Array.isArray(value)) {
              const existingStyle = p.style && typeof p.style === 'object' ? p.style : {};
              p.style = { ...existingStyle, ...(value as Record<string, unknown>) };
            } else {
              p[key] = value;
            }
          }
        });
        break;
      }

      case 'replace_node': {
        try {
          const nodeTree = buildNodeTree(query, op.tree);
          const parent: string = query.node(id).get().data.parent ?? 'ROOT';
          const siblings: string[] = query.node(parent).childNodes();
          const index = siblings.indexOf(id);
          actions.delete(id);
          actions.addNodeTree(nodeTree, parent, index);
        } catch (e) {
          console.warn('sitesmith patch: replace_node failed', e);
        }
        break;
      }

      case 'insert_after':
      case 'insert_before': {
        try {
          const nodeTree = buildNodeTree(query, op.tree);
          const parent: string = query.node(id).get().data.parent ?? 'ROOT';
          const siblings: string[] = query.node(parent).childNodes();
          const index = siblings.indexOf(id);
          const at = op.op === 'insert_after' ? index + 1 : index;
          actions.addNodeTree(nodeTree, parent, at);
        } catch (e) {
          console.warn(`sitesmith patch: ${op.op} failed`, e);
        }
        break;
      }

      case 'delete_node':
        try {
          actions.delete(id);
        } catch (e) {
          console.warn('sitesmith patch: delete_node failed', e);
        }
        break;

      default:
        console.warn('sitesmith patch: unknown op', (op as any).op);
    }
  }
  return { ok: true };
}