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

/** 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',
};

/**
 * 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 {
  const idCounter = { n: 0 };
  const craftNodes: Record<string, any> = {};

  const walk = (node: SerializedTreeNode, parent: string | null): string => {
    const id = (node.props.node_id as string | undefined) || `ai-auto-${idCounter.n++}`;
    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(tree, 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 };

/**
 * 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':
        actions.setProp(id, (p: any) => { Object.assign(p, op.props); });
        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 };
}