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

import { useEditor } from '@craftjs/core';
import { usePages } from '../state/PageContext';
import { SitesmithResponse, SerializedTreeNode, SitesmithPatchOp } from '../types/sitesmith';
import { sanitizeAiTree, buildNodeTree } from './craft-tree';

// Re-exported for existing callers/tests that import sanitizeAiTree/
// buildNodeTree from this module — the implementation lives in craft-tree.ts
// alongside the shared flattener: `templates/apply-template.ts` needs the
// exact same tree-materialization logic to load template header/footer/page
// content (see the doc comment on `buildNodeTree` in craft-tree.ts).
export { sanitizeAiTree, buildNodeTree };

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

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