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

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;
  /**
   * Bookkeeping-only: point `activePageId` at an already-loaded page without
   * re-serializing/deserializing the canvas (the caller -- e.g. useWhpApi's
   * `load()` -- has already put the right state on the canvas itself). Used
   * to fix I-1: `activePageId` defaults to the hardcoded `'home'` and
   * `load()` never updated it, so after the original Home page was deleted
   * (its replacement gets a fresh `page_<ts>` id) and the app reloaded,
   * `activePageId` pointed at nothing in `pages`, and a subsequent edit+save
   * only reached the legacy top-level fields.
   */
  setActivePageIdDirect: (pageId: string) => 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":{}}}';

/**
 * Flatten a `SerializedTreeNode` (as produced by the AI, a template, etc.)
 * into a Craft.js `SerializedNodes` JSON string ready for
 * `actions.deserialize()`.
 *
 * Untrusted input is validated the same way `apply-ai-response.ts`'s
 * `buildNodeTree` validates AI `patch`/section-replace trees: any node whose
 * `type.resolvedName` isn't a registered component is dropped (subtree and
 * all) via `sanitizeAiTree`, and a bad/colliding/`'ROOT'` id is regenerated
 * — never left as-is. This matters here specifically because `replace`
 * scope `site`/`page` responses reach this function via `actions.deserialize`
 * with no further validation downstream, unlike the `buildNodeTree` path
 * which also gets Craft.js's own `parseFreshNode` as a second line of
 * defense. If the ROOT node itself is invalid, sanitizeAiTree returns null
 * and we fall back to an empty canvas rather than handing deserialize()
 * something that could throw.
 *
 * Exported standalone (not a hook) so it's directly unit-testable without
 * mounting a Craft.js `<Editor>`.
 */
export function treeToCraftState(tree: SerializedTreeNode): string {
  const sanitized = sanitizeAiTree(tree, new Set());
  if (!sanitized) return EMPTY_CANVAS;

  const { rootNodeId, nodes: flatNodes } = flattenTreeForCraft(sanitized);
  const nodes: Record<string, FlatCraftNode> = flatNodes;

  // Craft.js deserialize requires the root node keyed as 'ROOT'.
  if (rootNodeId !== 'ROOT') {
    nodes['ROOT'] = { ...nodes[rootNodeId], parent: null, isCanvas: true };
    delete nodes[rootNodeId];
    for (const childId of nodes['ROOT'].nodes) {
      if (nodes[childId]) nodes[childId].parent = 'ROOT';
    }
  }
  return JSON.stringify(nodes);
}

// Default header seed: a ROOT header Container holding a Navbar with default
// links. New sites previously opened with an EMPTY header (just a bare
// Container), so there was no menu to edit and the empty zone rendered as a
// stray band above the page. Seeding a real Navbar gives every new site an
// editable menu-with-links out of the box (and removes the empty-header gap).
// Node shape matches treeToCraftState() / Craft's actions.deserialize().
export const DEFAULT_HEADER_STATE = JSON.stringify({
  ROOT: {
    type: { resolvedName: 'Container' },
    isCanvas: true,
    props: { style: { width: '100%' }, tag: 'header' },
    displayName: 'Container',
    custom: {},
    hidden: false,
    nodes: ['header-navbar'],
    linkedNodes: {},
  },
  'header-navbar': {
    type: { resolvedName: 'Navbar' },
    isCanvas: false,
    props: {
      logoType: 'text',
      logoText: 'MySite',
      logoImage: '',
      logoWidth: '120px',
      logoUrl: '/',
      logoFontFamily: 'Inter, sans-serif',
      logoFontSize: '20px',
      links: [
        { text: 'Home', href: '/' },
        { text: 'About', href: '#about' },
        { text: 'Services', href: '#services' },
        { text: 'Contact', href: '#contact', isCta: true },
      ],
      backgroundColor: '#ffffff',
      textColor: '#3f3f46',
      hoverColor: '#3b82f6',
      ctaColor: '#3b82f6',
      ctaTextColor: '#ffffff',
      padding: '16px 24px',
      navAlignment: 'space-between',
      isSticky: false,
      showMobileMenu: false,
      style: { borderBottom: '1px solid #e4e4e7' },
    },
    displayName: 'Navbar',
    parent: 'ROOT',
    custom: {},
    hidden: false,
    nodes: [],
    linkedNodes: {},
  },
});

const PageContext = createContext<PageContextValue>({
  pages: [],
  headerPage: { id: HEADER_ID, name: 'Header', slug: '__header__', craftState: null },
  footerPage: { id: FOOTER_ID, name: 'Footer', slug: '__footer__', craftState: null },
  activePageId: 'home',
  isEditingHeader: false,
  isEditingFooter: false,
  switchPage: () => {},
  editHeader: () => {},
  editFooter: () => {},
  addPage: () => {},
  deletePage: () => {},
  renamePage: () => {},
  setHeaderCraftState: () => {},
  setFooterCraftState: () => {},
  setPagesCraftState: () => {},
  setActivePageIdDirect: () => {},
  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, '-');
}

/**
 * Append `-2`, `-3`, … to `base` until it no longer collides with
 * `existingSlugs`. Two pages that slugify to the same string (e.g. both
 * named "About") must not both publish to `about.html` — the second write
 * would silently overwrite the first on publish.
 */
export function uniqueSlug(base: string, existingSlugs: string[]): string {
  if (!existingSlugs.includes(base)) return base;
  let i = 2;
  while (existingSlugs.includes(`${base}-${i}`)) i++;
  return `${base}-${i}`;
}

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

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

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

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;

  // Mirror the latest state in refs so event handlers (switchPage,
  // deletePage) can synchronously read "what's current" without smuggling a
  // side effect into a setState updater function to peek at `prev` — updater
  // functions must stay pure since React (StrictMode) may invoke them twice.
  const pagesRef = useRef(pages);
  pagesRef.current = pages;
  const headerPageRef = useRef(headerPage);
  headerPageRef.current = headerPage;
  const footerPageRef = useRef(footerPage);
  footerPageRef.current = footerPage;

  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. `pageId` can never equal `currentId` here
      // (guarded by the early return above), so the target's stored state
      // was untouched by the setState calls just above — the refs (kept in
      // sync with state on every render) are safe to read synchronously
      // without waiting for a re-render, and without running the load as a
      // side effect inside a setState updater.
      if (pageId === HEADER_ID) {
        loadState(headerPageRef.current.craftState, EMPTY_HEADER);
      } else if (pageId === FOOTER_ID) {
        loadState(footerPageRef.current.craftState, EMPTY_FOOTER);
      } else {
        const target = pagesRef.current.find((p) => p.id === pageId);
        loadState(target?.craftState || null, EMPTY_CANVAS);
      }

      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 requestedSlug = slug || slugify(name);
      const id = `page_${Date.now()}`;

      // Save current page first
      saveCurrentState();

      setPages((prev) => [
        ...prev,
        {
          id,
          name,
          slug: uniqueSlug(requestedSlug, prev.map((p) => p.slug)),
          craftState: null,
        },
      ]);

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

      const prev = pagesRef.current;
      if (prev.length <= 1) return;

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

      // If deleting the active page, switch to the first remaining. This
      // runs after the state update is computed (not inside the setPages
      // updater) so the side effects (deserialize, ref/state mutation) fire
      // exactly once regardless of how many times React invokes updaters.
      if (pageId === activePageIdRef.current) {
        const nextPage = filtered[0];
        setActivePageId(nextPage.id);
        activePageIdRef.current = nextPage.id;
        loadState(nextPage.craftState, EMPTY_CANVAS);
      }
    },
    [loadState],
  );

  const renamePage = useCallback((pageId: string, name: string, slug: string) => {
    setPages((prev) =>
      prev.map((p, i) => {
        if (p.id !== pageId) return p;
        // First page is the landing page — its slug is locked to 'index' so
        // the file always publishes to index.html regardless of the user-set
        // name. The display name can change freely.
        if (i === 0) return { ...p, name, slug: 'index' };
        const requestedSlug = slug || slugify(name);
        const otherSlugs = prev.filter((pp) => pp.id !== pageId).map((pp) => pp.slug);
        return { ...p, name, slug: uniqueSlug(requestedSlug, otherSlugs) };
      }),
    );
  }, []);

  /** 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 }));
  }, []);

  /**
   * Bookkeeping-only setter for `activePageId` -- see the doc comment on
   * `PageContextValue.setActivePageIdDirect`. Does NOT serialize/deserialize
   * the canvas; callers that need that should use `switchPage` instead.
   */
  const setActivePageIdDirect = useCallback((pageId: string) => {
    setActivePageId(pageId);
    activePageIdRef.current = pageId;
  }, []);

  /** 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, i) => ({
      id: p.id,
      name: p.name,
      // Heal legacy projects whose first page was saved with slug='home' (or
      // any other) before the landing-page rule existed. The first page is
      // ALWAYS the landing page → slug 'index' → file index.html.
      slug: i === 0 ? 'index' : p.slug,
      craftState: p.craftState,
    })));
  }, []);

  /**
   * 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;
    // Track slugs as they're assigned so later pages dedupe against earlier
    // ones in the same batch (e.g. the AI generating two "About" pages).
    const seenSlugs: string[] = [];
    const built = newPages.map((p, i) => {
      // First page must publish to index.html so it serves at the site root.
      // Apache resolves '/' to index.html, not home.html — without this, the
      // AI's "Home" page lands at /home.html and visitors hit a blank root.
      const slug = i === 0 ? 'index' : uniqueSlug(slugify(p.name), seenSlugs);
      seenSlugs.push(slug);
      return {
        id: i === 0 ? 'home' : `page_${Date.now()}_${i}`,
        name: p.name,
        slug,
        craftState: treeToCraftState(p.tree),
      };
    });
    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);
  }, [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 = treeToCraftState(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);
  }, [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 = treeToCraftState(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);
    }
  }, [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 = treeToCraftState(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);
    }
  }, [loadState]);

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