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

/** Component types that act as drop targets (isCanvas = true) */
const CANVAS_TYPES = new Set([
  'Container', 'Section', 'ColumnLayout', 'BackgroundSection',
  'HeroSimple', 'FeaturesGrid', 'CTASection',
  'FormContainer', 'Navbar', 'Footer',
]);

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

  return async function apply(resp: SitesmithResponse): 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') {
        // 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 };
}