site-builder/craft/src/state/PageContext.tsx

import React, { createContext, useContext, useState, useCallback, useRef, ReactNode } from 'react';
import { useEditor } from '@craftjs/core';
import { PageData } from '../types';
import { SerializedTreeNode } from '../types/sitesmith';
import { useSiteDesign, SiteDesign } from './SiteDesignContext';

/** Only `Container` instances are "real" canvases in serialized state — they
 * directly render whatever is in node.data.nodes. Layout-shell components
 * (Section, HeroSimple, FeaturesGrid, ColumnLayout, CTASection, etc) use
 * Craft.js <Element canvas id="…"> linkedNodes internally; their own
 * isCanvas must be FALSE or Craft.js's toNodeTree walker trips an Invariant
 * because the shell claims to be a canvas but its render ignores `nodes`. */
const CANVAS_TYPES = new Set<string>(['Container']);

/** Shells that wrap their content in a single <Element id="<key>" is={Container}>.
 * When the AI puts content directly under one of these, the children end up
 * orphaned (the shell ignores data.nodes — it renders via the linkedNode) and
 * Craft.js auto-creates the linkedNode at render time with a botched type
 * field, which then crashes toNodeTree. Pre-create the linkedNode ourselves
 * to keep the state shape Craft.js expects. */
const SHELL_INNER: Record<string, string> = {
  Section: 'section-inner',
  BackgroundSection: 'bg-section-inner',
  FormContainer: 'form-inner',
};

interface PageContextValue {
  pages: PageData[];
  headerPage: PageData;
  footerPage: PageData;
  activePageId: string;
  isEditingHeader: boolean;
  isEditingFooter: boolean;
  switchPage: (pageId: string) => void;
  editHeader: () => void;
  editFooter: () => void;
  addPage: (name: string, slug: string) => void;
  deletePage: (pageId: string) => void;
  renamePage: (pageId: string, name: string, slug: string) => void;
  setHeaderCraftState: (craftState: string) => void;
  setFooterCraftState: (craftState: string) => void;
  setPagesCraftState: (pagesData: { id: string; name: string; slug: string; craftState: string | null }[]) => void;
  /** AI helpers — replace entire site or page with a new tree */
  replaceAllPages: (pages: { name: string; tree: SerializedTreeNode }[]) => void;
  replaceCurrentPage: (page: { name: string; tree: SerializedTreeNode }) => void;
  setHeader: (tree: SerializedTreeNode) => void;
  setFooter: (tree: SerializedTreeNode) => void;
  siteDesign: SiteDesign;
}

const HEADER_ID = '__header__';
const FOOTER_ID = '__footer__';

const EMPTY_CANVAS =
  '{"ROOT":{"type":{"resolvedName":"Container"},"isCanvas":true,"props":{"style":{"minHeight":"100vh","backgroundColor":"#ffffff"},"tag":"div"},"displayName":"Container","custom":{},"hidden":false,"nodes":[],"linkedNodes":{}}}';

const EMPTY_HEADER =
  '{"ROOT":{"type":{"resolvedName":"Container"},"isCanvas":true,"props":{"style":{"minHeight":"60px","backgroundColor":"#ffffff","padding":"12px 24px","display":"flex","alignItems":"center"},"tag":"header"},"displayName":"Container","custom":{},"hidden":false,"nodes":[],"linkedNodes":{}}}';

const EMPTY_FOOTER =
  '{"ROOT":{"type":{"resolvedName":"Container"},"isCanvas":true,"props":{"style":{"minHeight":"60px","backgroundColor":"#0f172a","color":"#94a3b8","padding":"40px 24px","textAlign":"center"},"tag":"footer"},"displayName":"Container","custom":{},"hidden":false,"nodes":[],"linkedNodes":{}}}';

const PageContext = createContext<PageContextValue>({
  pages: [],
  headerPage: { id: HEADER_ID, name: 'Header', slug: '__header__', craftState: null, headCode: '' },
  footerPage: { id: FOOTER_ID, name: 'Footer', slug: '__footer__', craftState: null, headCode: '' },
  activePageId: 'home',
  isEditingHeader: false,
  isEditingFooter: false,
  switchPage: () => {},
  editHeader: () => {},
  editFooter: () => {},
  addPage: () => {},
  deletePage: () => {},
  renamePage: () => {},
  setHeaderCraftState: () => {},
  setFooterCraftState: () => {},
  setPagesCraftState: () => {},
  replaceAllPages: () => {},
  replaceCurrentPage: () => {},
  setHeader: () => {},
  setFooter: () => {},
  siteDesign: {} as SiteDesign,
});

export const usePages = () => useContext(PageContext);

function slugify(name: string): string {
  return name
    .toLowerCase()
    .trim()
    .replace(/[^a-z0-9\s-]/g, '')
    .replace(/\s+/g, '-')
    .replace(/-+/g, '-');
}

const DEFAULT_PAGE: PageData = {
  id: 'home',
  name: 'Home',
  slug: 'index',
  craftState: null,
  headCode: '',
};

const DEFAULT_HEADER: PageData = {
  id: HEADER_ID,
  name: 'Header',
  slug: '__header__',
  craftState: null,
  headCode: '',
};

const DEFAULT_FOOTER: PageData = {
  id: FOOTER_ID,
  name: 'Footer',
  slug: '__footer__',
  craftState: null,
  headCode: '',
};

export const PageProvider: React.FC<{ children: ReactNode }> = ({ children }) => {
  const { query, actions } = useEditor();
  const { design } = useSiteDesign();
  const [pages, setPages] = useState<PageData[]>([DEFAULT_PAGE]);
  const [headerPage, setHeaderPage] = useState<PageData>(DEFAULT_HEADER);
  const [footerPage, setFooterPage] = useState<PageData>(DEFAULT_FOOTER);
  const [activePageId, setActivePageId] = useState('home');
  const activePageIdRef = useRef(activePageId);
  activePageIdRef.current = activePageId;

  const isEditingHeader = activePageId === HEADER_ID;
  const isEditingFooter = activePageId === FOOTER_ID;

  /** Save whatever is on the current Frame back to the right state slot */
  const saveCurrentState = useCallback(() => {
    const currentState = query.serialize();
    const currentId = activePageIdRef.current;

    if (currentId === HEADER_ID) {
      setHeaderPage((prev) => ({ ...prev, craftState: currentState }));
    } else if (currentId === FOOTER_ID) {
      setFooterPage((prev) => ({ ...prev, craftState: currentState }));
    } else {
      setPages((prev) =>
        prev.map((p) => (p.id === currentId ? { ...p, craftState: currentState } : p)),
      );
    }
  }, [query]);

  /** Load a craft state into the Frame */
  const loadState = useCallback(
    (craftState: string | null, fallback: string) => {
      setTimeout(() => {
        try {
          actions.deserialize(craftState || fallback);
        } catch (e) {
          console.error('Failed to deserialize state:', e);
          try {
            actions.deserialize(fallback);
          } catch (_e2) {
            // give up
          }
        }
      }, 0);
    },
    [actions],
  );

  const switchPage = useCallback(
    (pageId: string) => {
      if (pageId === activePageIdRef.current) return;

      // Serialize the current Craft.js state synchronously BEFORE switching
      const currentState = query.serialize();
      const currentId = activePageIdRef.current;

      // Persist the serialized state to the correct page slot
      if (currentId === HEADER_ID) {
        setHeaderPage((prev) => ({ ...prev, craftState: currentState }));
      } else if (currentId === FOOTER_ID) {
        setFooterPage((prev) => ({ ...prev, craftState: currentState }));
      } else {
        setPages((prev) =>
          prev.map((p) => (p.id === currentId ? { ...p, craftState: currentState } : p)),
        );
      }

      // Load target page state.
      // For header/footer we need the latest saved state. Since setState above is async,
      // read from the ref-like state getter. For header/footer, we read the current
      // state value and fall back to what we just saved if the target is the same slot.
      if (pageId === HEADER_ID) {
        // Use functional state read to get the latest value
        setHeaderPage((prev) => {
          loadState(prev.craftState, EMPTY_HEADER);
          return prev;
        });
      } else if (pageId === FOOTER_ID) {
        setFooterPage((prev) => {
          loadState(prev.craftState, EMPTY_FOOTER);
          return prev;
        });
      } else {
        setPages((prev) => {
          const target = prev.find((p) => p.id === pageId);
          loadState(target?.craftState || null, EMPTY_CANVAS);
          return prev;
        });
      }

      setActivePageId(pageId);
      activePageIdRef.current = pageId;
    },
    [query, loadState],
  );

  const editHeader = useCallback(() => {
    switchPage(HEADER_ID);
  }, [switchPage]);

  const editFooter = useCallback(() => {
    switchPage(FOOTER_ID);
  }, [switchPage]);

  const addPage = useCallback(
    (name: string, slug: string) => {
      const finalSlug = slug || slugify(name);
      const id = `page_${Date.now()}`;

      // Save current page first
      saveCurrentState();

      setPages((prev) => [
        ...prev,
        {
          id,
          name,
          slug: finalSlug,
          craftState: null,
          headCode: '',
        },
      ]);

      // Switch to the new page with empty canvas
      loadState(null, EMPTY_CANVAS);

      setActivePageId(id);
      activePageIdRef.current = id;
    },
    [saveCurrentState, loadState],
  );

  const deletePage = useCallback(
    (pageId: string) => {
      // Can't delete header/footer
      if (pageId === HEADER_ID || pageId === FOOTER_ID) return;

      setPages((prev) => {
        if (prev.length <= 1) return prev;

        const filtered = prev.filter((p) => p.id !== pageId);

        // If deleting the active page, switch to the first remaining
        if (pageId === activePageIdRef.current) {
          const nextPage = filtered[0];
          setActivePageId(nextPage.id);
          activePageIdRef.current = nextPage.id;
          loadState(nextPage.craftState, EMPTY_CANVAS);
        }

        return filtered;
      });
    },
    [loadState],
  );

  const renamePage = useCallback((pageId: string, name: string, slug: string) => {
    setPages((prev) =>
      prev.map((p) =>
        p.id === pageId ? { ...p, name, slug: slug || slugify(name) } : p,
      ),
    );
  }, []);

  /** Allow external code (e.g., load from API) to set the header craft state */
  const setHeaderCraftState = useCallback((craftState: string) => {
    setHeaderPage((prev) => ({ ...prev, craftState }));
  }, []);

  /** Allow external code (e.g., load from API) to set the footer craft state */
  const setFooterCraftState = useCallback((craftState: string) => {
    setFooterPage((prev) => ({ ...prev, craftState }));
  }, []);

  /** Allow external code (e.g., load from API) to restore pages with craft states */
  const setPagesCraftState = useCallback((pagesData: { id: string; name: string; slug: string; craftState: string | null }[]) => {
    setPages(pagesData.map((p) => ({
      id: p.id,
      name: p.name,
      slug: p.slug,
      craftState: p.craftState,
      headCode: '',
    })));
  }, []);

  /** Flatten a SerializedTreeNode into a Craft.js SerializedNodes JSON string */
  const treeToState = useCallback((tree: SerializedTreeNode): string => {
    let counter = 0;
    const nodes: Record<string, unknown> = {};
    const walk = (node: SerializedTreeNode, parent: string | null): string => {
      const id = (node.props.node_id as string | undefined) || `ai-auto-${counter++}`;
      const childIds: string[] = [];
      const typeName = node.type?.resolvedName;
      // Normalize props: the AI sometimes emits `style: []` instead of `{}`.
      // React/Craft.js choke when a CSSProperties slot is an array — normalize it.
      const rawProps = node.props ?? {};
      const props: Record<string, unknown> = { ...rawProps };
      if (Array.isArray(props.style)) props.style = {};
      nodes[id] = {
        type: node.type,
        // isCanvas must match the component's craft.rules — only layout
        // wrappers accept children. Setting it true on leaf components
        // (Heading, TextBlock, ButtonLink, etc) makes Craft.js render them
        // as empty drop-canvas wrappers and the actual content disappears.
        isCanvas: typeName ? CANVAS_TYPES.has(typeName) : false,
        props,
        displayName: typeName,
        custom: {},
        hidden: false,
        parent,
        nodes: childIds,
        linkedNodes: {},
      };
      for (const child of node.nodes ?? []) {
        childIds.push(walk(child, id));
      }
      // ColumnLayout uses Craft.js linkedNodes with fixed ids (col-0, col-1, ...).
      // The AI emits children as direct `nodes`, but ColumnLayout's render ignores
      // them and creates fresh column Elements — the AI's children become orphans
      // and any subsequent toNodeTree walk hits an Invariant. Move direct children
      // into linkedNodes so they render in the columns the user actually sees.
      if (typeName === 'ColumnLayout' && childIds.length > 0) {
        const linked: Record<string, string> = {};
        childIds.forEach((cid, i) => {
          linked[`col-${i}`] = cid;
          if (nodes[cid]) (nodes[cid] as any).isCanvas = true; // columns are canvases
        });
        (nodes[id] as any).nodes = [];
        (nodes[id] as any).linkedNodes = linked;
        // Reflect the actual column count on the component so its render matches.
        const cur = (nodes[id] as any).props || {};
        if (!cur.columns || cur.columns !== childIds.length) cur.columns = childIds.length;
        (nodes[id] as any).props = cur;
      }
      // Section/BackgroundSection/FormContainer each render a single
      // <Element id="<key>" is={Container} canvas> ... </Element>. If the AI
      // nests content as direct children, Craft.js will auto-create the
      // linkedNode on first render — and store its type as the Container
      // component class rather than {resolvedName:'Container'}, which then
      // crashes toNodeTree with "type (undefined) does not exist in resolver".
      // Pre-create the linkedNode ourselves with the correct serialized type
      // so Craft.js never has to materialize it.
      const innerKey = SHELL_INNER[typeName ?? ''];
      if (innerKey && childIds.length > 0) {
        const innerId = `${id}__${innerKey}`;
        nodes[innerId] = {
          type: { resolvedName: 'Container' },
          isCanvas: true,
          props: { tag: 'div' },
          displayName: 'Container',
          custom: {},
          hidden: false,
          parent: id,
          nodes: [...childIds],
          linkedNodes: {},
        };
        for (const cid of childIds) {
          if (nodes[cid]) (nodes[cid] as any).parent = innerId;
        }
        (nodes[id] as any).nodes = [];
        (nodes[id] as any).linkedNodes = { [innerKey]: innerId };
      }
      return id;
    };
    const rootId = walk(tree, null);
    // Craft.js deserialize requires the root node keyed as 'ROOT'
    if (rootId !== 'ROOT') {
      nodes['ROOT'] = nodes[rootId];
      (nodes['ROOT'] as any).parent = null;
      // ROOT must be a canvas regardless of component type so children render.
      (nodes['ROOT'] as any).isCanvas = true;
      delete nodes[rootId];
      // Fix up parent references from ROOT's children
      for (const childId of (nodes['ROOT'] as any).nodes) {
        if (nodes[childId]) (nodes[childId] as any).parent = 'ROOT';
      }
    }
    return JSON.stringify(nodes);
  }, []);

  /**
   * AI helper: replace all pages with newly generated trees.
   * Stores each page's serialized state without touching the live canvas
   * (the canvas still shows the currently active page — call switchPage() if needed).
   */
  const replaceAllPages = useCallback((newPages: { name: string; tree: SerializedTreeNode }[]) => {
    if (newPages.length === 0) return;
    const built = newPages.map((p, i) => ({
      id: i === 0 ? 'home' : `page_${Date.now()}_${i}`,
      name: p.name,
      slug: slugify(p.name),
      craftState: treeToState(p.tree),
      headCode: '',
    }));
    setPages(built);
    // Load the first page into the live canvas
    const firstState = built[0].craftState;
    setActivePageId(built[0].id);
    activePageIdRef.current = built[0].id;
    loadState(firstState, EMPTY_CANVAS);
  }, [treeToState, loadState]);

  /**
   * AI helper: replace the current page's tree.
   * Deserializes the new tree into the live Craft.js canvas and persists it.
   */
  const replaceCurrentPage = useCallback((page: { name: string; tree: SerializedTreeNode }) => {
    const craftState = treeToState(page.tree);
    const currentId = activePageIdRef.current;
    if (currentId === HEADER_ID) {
      setHeaderPage((prev) => ({ ...prev, name: page.name, craftState }));
    } else if (currentId === FOOTER_ID) {
      setFooterPage((prev) => ({ ...prev, name: page.name, craftState }));
    } else {
      setPages((prev) =>
        prev.map((p) => (p.id === currentId ? { ...p, name: page.name, craftState } : p)),
      );
    }
    loadState(craftState, EMPTY_CANVAS);
  }, [treeToState, loadState]);

  /**
   * AI helper: replace the shared header tree.
   * Updates stored state; does NOT switch the canvas to header view.
   */
  const setHeader = useCallback((tree: SerializedTreeNode) => {
    const craftState = treeToState(tree);
    setHeaderPage((prev) => ({ ...prev, craftState }));
    // If the canvas is currently showing the header, refresh it live
    if (activePageIdRef.current === HEADER_ID) {
      loadState(craftState, EMPTY_HEADER);
    }
  }, [treeToState, loadState]);

  /**
   * AI helper: replace the shared footer tree.
   * Updates stored state; does NOT switch the canvas to footer view.
   */
  const setFooter = useCallback((tree: SerializedTreeNode) => {
    const craftState = treeToState(tree);
    setFooterPage((prev) => ({ ...prev, craftState }));
    // If the canvas is currently showing the footer, refresh it live
    if (activePageIdRef.current === FOOTER_ID) {
      loadState(craftState, EMPTY_FOOTER);
    }
  }, [treeToState, loadState]);

  return (
    <PageContext.Provider
      value={{
        pages,
        headerPage,
        footerPage,
        activePageId,
        isEditingHeader,
        isEditingFooter,
        switchPage,
        editHeader,
        editFooter,
        addPage,
        deletePage,
        renamePage,
        setHeaderCraftState,
        setFooterCraftState,
        setPagesCraftState,
        replaceAllPages,
        replaceCurrentPage,
        setHeader,
        setFooter,
        siteDesign: design,
      }}
    >
      {children}
    </PageContext.Provider>
  );
};