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, SitesmithPatchOp } from '../types/sitesmith';
import { sanitizeAiTree, flattenTreeForCraft } from './craft-tree';

// Re-exported for existing callers/tests that import sanitizeAiTree from
// this module — the implementation now lives in craft-tree.ts alongside the
// shared flattener, since treeToState (PageContext.tsx) needs the same
// validation this module pioneered.
export { sanitizeAiTree };

/**
 * 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']);

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

  // sanitizeAiTree() has already guaranteed every node has a valid, unique,
  // non-'ROOT' node_id and a registered resolvedName — the shared flattener
  // 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)) {
    if (flat.syntheticIds.has(id)) {
      // The SHELL_INNER wrapper Container is synthesized by the flattener,
      // not present in the AI's 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.
      craftNodes[id] = {
        id,
        data: { ...flatNode },
        events: { selected: false, hovered: false, dragged: false },
        rules: { canDrag: () => true, canMoveIn: () => true, canMoveOut: () => true, canDrop: () => true },
      };
    } else {
      craftNodes[id] = (query.parseFreshNode({ id, data: { ...flatNode } }) as any).toNode();
    }
  }

  return { rootNodeId: flat.rootNodeId, 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: SitesmithPatchOp[],
): { 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') {
              if (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 {
                // `style: null` (or any other non-object) is not a valid
                // style patch — ignore it rather than clobbering the whole
                // style object with null/undefined/a stray primitive.
                console.warn('sitesmith patch: update_props ignoring non-object "style" value', value);
              }
            } 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:
        // Every SitesmithPatchOp variant is handled above, so `op` is typed
        // `never` here at compile time — but the AI response is untrusted
        // JSON (see useSitesmith.ts's `j: SendResult = await r.json()`), so
        // a malformed/unrecognized `op` field can still reach this branch at
        // runtime. Log the raw value rather than assuming a shape.
        console.warn('sitesmith patch: unknown op', op);
    }
  }
  return { ok: true };
}