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

import React, { createContext, useCallback, useContext, useEffect, useMemo, useRef, useState, ReactNode } from 'react';
import { useEditorConfig } from './EditorConfigContext';

export type LoadState = 'loading' | 'loaded' | 'failed';

interface SaveGuardContextValue {
  /** Nothing may be saved until this is 'loaded' -- a failed load leaves the
   *  canvas blank, and saving that would overwrite the customer's site. */
  loadState: LoadState;
  setLoadState: (s: LoadState) => void;
  /** Server refused a save that would replace a non-empty site with an empty
   *  one (409 empty_overwrite_refused). Holds the message while the banner is up. */
  emptyRefusal: string | null;
  setEmptyRefusal: (message: string | null) => void;
  /** One-shot: the next save sends force_empty (consumed by the next successful save). */
  forceEmptyNextSave: React.MutableRefObject<boolean>;
  armForceEmptyNextSave: () => void;
  /** True when the editor holds changes (canvas, pages, design, header/footer)
   *  made after the last successful save or load. Cleared only by a save. */
  dirty: boolean;
  /** Synchronous read of `dirty`, for handlers that must not see a stale render. */
  isDirty: () => boolean;
  /** Record a change. Ignored until the site has loaded (load itself rewrites state). */
  markDirty: () => void;
  /** Monotonic change counter: compare two reads to learn whether anything changed in between. */
  getVersion: () => number;
  /** Call right before serializing for a save; pass the ticket to `endSave`. */
  beginSave: () => SaveTicket;
  /** Finish a save. Clears dirty only if it succeeded AND nothing changed since `beginSave`. */
  endSave: (ticket: SaveTicket, succeeded: boolean) => void;
  /** True while at least one save request is in flight. */
  saving: boolean;
  /** Skip the unsaved-changes prompt for the next unload (e.g. a deliberate "Reload latest"). */
  suppressBeforeUnload: () => void;
  /** Server revision the editor last loaded or saved: a string, null (no project on the
   *  server yet), or undefined (unknown -- an older backend sent none, so saves skip the check). */
  getRev: () => string | null | undefined;
  setRev: (rev: string | null) => void;
  /** Server message while the "changed in another tab" modal is up; autosave stays paused. */
  conflict: string | null;
  setConflict: (message: string | null) => void;
  /** One-shot: the next save sends force_overwrite (armed by "Keep mine", spent by a save that goes through). */
  forceOverwriteNextSave: React.MutableRefObject<boolean>;
  armForceOverwriteNextSave: () => void;
  /** Cancel an armed overwrite (the customer backed out of the follow-up empty-save prompt). */
  disarmForceOverwriteNextSave: () => void;
  /** "Reload latest": skip the unsaved-changes prompt, then reload the whole editor. */
  reloadLatest: () => void;
}

/** Identifies the editor state a save was built from. */
export interface SaveTicket {
  version: number;
}

// Default (no provider: standalone mode and isolated unit tests) is
// permissive -- the gate only applies inside <SaveGuardProvider> in WHP mode.
const defaultValue: SaveGuardContextValue = {
  loadState: 'loaded',
  setLoadState: () => {},
  emptyRefusal: null,
  setEmptyRefusal: () => {},
  forceEmptyNextSave: { current: false },
  armForceEmptyNextSave: () => {},
  dirty: false,
  isDirty: () => false,
  getVersion: () => 0,
  markDirty: () => {},
  beginSave: () => ({ version: 0 }),
  endSave: () => {},
  saving: false,
  suppressBeforeUnload: () => {},
  getRev: () => undefined,
  setRev: () => {},
  conflict: null,
  setConflict: () => {},
  forceOverwriteNextSave: { current: false },
  armForceOverwriteNextSave: () => {},
  disarmForceOverwriteNextSave: () => {},
  reloadLatest: () => {},
};

const SaveGuardContext = createContext<SaveGuardContextValue>(defaultValue);

export const useSaveGuard = () => useContext(SaveGuardContext);

export const SaveGuardProvider: React.FC<{ children: ReactNode; /** Injectable for tests; defaults to a full page reload. */ reloadPage?: () => void }> = ({ children, reloadPage }) => {
  const { isWHP } = useEditorConfig();
  const [loadState, setLoadState] = useState<LoadState>(isWHP ? 'loading' : 'loaded');
  const [emptyRefusal, setEmptyRefusal] = useState<string | null>(null);
  const forceEmptyNextSave = useRef(false);
  const [conflict, setConflict] = useState<string | null>(null);
  const forceOverwriteNextSave = useRef(false);
  const revRef = useRef<string | null | undefined>(undefined);
  const getRev = useCallback(() => revRef.current, []);
  const setRev = useCallback((rev: string | null) => { revRef.current = rev; }, []);

  // Dirty tracking. `versionRef` counts changes so a save that raced with an
  // edit cannot mark that edit clean.
  const [dirty, setDirty] = useState(false);
  const [savesInFlight, setSavesInFlight] = useState(0);
  const dirtyRef = useRef(false);
  const versionRef = useRef(0);
  const loadedRef = useRef(loadState === 'loaded');
  const suppressRef = useRef(false);

  // A (re)load replaces the whole editor state: it is the new clean baseline.
  useEffect(() => {
    loadedRef.current = loadState === 'loaded';
    if (loadState === 'loaded') {
      dirtyRef.current = false;
      setDirty(false);
    }
  }, [loadState]);

  const markDirty = useCallback(() => {
    if (!loadedRef.current) return;
    versionRef.current += 1;
    dirtyRef.current = true;
    setDirty(true);
  }, []);
  const isDirty = useCallback(() => dirtyRef.current, []);
  const getVersion = useCallback(() => versionRef.current, []);
  const beginSave = useCallback((): SaveTicket => {
    setSavesInFlight((n) => n + 1);
    return { version: versionRef.current };
  }, []);
  const endSave = useCallback((ticket: SaveTicket, succeeded: boolean) => {
    setSavesInFlight((n) => Math.max(0, n - 1));
    if (succeeded && ticket.version === versionRef.current) {
      dirtyRef.current = false;
      setDirty(false);
    }
  }, []);
  const suppressBeforeUnload = useCallback(() => { suppressRef.current = true; }, []);

  const reloadLatest = useCallback(() => {
    suppressRef.current = true;
    (reloadPage ?? (() => window.location.reload()))();
  }, [reloadPage]);

  const saving = savesInFlight > 0;
  const needsPrompt = isWHP && (dirty || saving);
  useEffect(() => {
    if (!needsPrompt) return;
    const onBeforeUnload = (e: BeforeUnloadEvent) => {
      if (suppressRef.current) return;
      e.preventDefault();
      e.returnValue = '';
    };
    window.addEventListener('beforeunload', onBeforeUnload);
    return () => window.removeEventListener('beforeunload', onBeforeUnload);
  }, [needsPrompt]);

  const value = useMemo<SaveGuardContextValue>(() => ({
    loadState,
    setLoadState,
    emptyRefusal,
    setEmptyRefusal,
    forceEmptyNextSave,
    armForceEmptyNextSave: () => { forceEmptyNextSave.current = true; },
    dirty,
    isDirty,
    getVersion,
    markDirty,
    beginSave,
    endSave,
    saving,
    suppressBeforeUnload,
    getRev,
    setRev,
    conflict,
    setConflict,
    forceOverwriteNextSave,
    armForceOverwriteNextSave: () => { forceOverwriteNextSave.current = true; },
    disarmForceOverwriteNextSave: () => { forceOverwriteNextSave.current = false; },
    reloadLatest,
  }), [loadState, emptyRefusal, dirty, saving, isDirty, getVersion, markDirty, beginSave, endSave, suppressBeforeUnload, getRev, setRev, conflict, reloadLatest]);

  return <SaveGuardContext.Provider value={value}>{children}</SaveGuardContext.Provider>;
};