siftlode/frontend/src/components/ui/dialogFocus.ts

import { useEffect, useEffectEvent, useRef, useState, type RefObject } from "react";

// The focus half of the dialog contract, in ONE place. `ui/Modal` owns it for every dialog that goes
// through the shell; the onboarding wizard renders its own shell (a bare `ui/Overlay`, because it is
// not a titled dialog with a header row) and used to have NONE of this — no trap, no initial focus,
// no restore — so a keyboard user tabbed straight out of the first-login wizard into the app behind
// the scrim, where the focus ring is invisible. A second hand-written copy is
// how two shells drift apart; this is the shared one.

// Tabbable elements inside the dialog, in DOM order. `:not([tabindex="-1"])` drops programmatically
// focusable-only nodes; the visibility filter (offsetParent) drops display:none ones (e.g. a
// collapsed disclosure) so the trap never lands focus on something the user can't see.
const FOCUSABLE =
  'a[href], button:not([disabled]), textarea:not([disabled]), input:not([disabled]), select:not([disabled]), [tabindex]:not([tabindex="-1"])';

export const focusablesIn = (root: HTMLElement): HTMLElement[] =>
  [...root.querySelectorAll<HTMLElement>(FOCUSABLE)].filter((el) => el.offsetParent !== null);

/** A form's first field — where a dialog opened to fill something in takes the focus: its first
 *  control, unless one is marked `data-initial-focus`. The mark is for a form with a control above
 *  its fields that is not one of them — the issue composer's Edit / Preview switch took the focus,
 *  and typing went nowhere. */
export const firstFieldIn = (root: HTMLElement): HTMLElement | undefined => {
  const all = focusablesIn(root);
  return all.find((el) => el.hasAttribute("data-initial-focus")) ?? all[0];
};

/** Keep Tab / Shift+Tab cycling inside `scope` instead of reaching the inert page behind the
 *  backdrop. Call from a keydown handler that has already decided this dialog is the topmost layer —
 *  a nested dialog must defer to its child. `container` is the scope's own tabIndex=-1 panel, if it
 *  has one: focus can land there by clicking dialog dead space, and `Node.contains(self)` is true, so
 *  without naming it a Shift+Tab from the panel would leak to the background. */
export function trapTab(
  e: KeyboardEvent,
  scope: HTMLElement,
  container?: HTMLElement | null,
): void {
  const items = focusablesIn(scope);
  const first = items[0];
  const last = items[items.length - 1];
  if (!first || !last) {
    e.preventDefault();
    return;
  }
  const active = document.activeElement;
  if (!scope.contains(active) || (container != null && active === container)) {
    e.preventDefault();
    (e.shiftKey ? last : first).focus();
  } else if (e.shiftKey && active === first) {
    e.preventDefault();
    last.focus();
  } else if (!e.shiftKey && active === last) {
    e.preventDefault();
    first.focus();
  }
}

// ONE body scroll lock for every open dialog, counted. Each dialog used to save and restore
// `overflow` on its own, so two unmounting out of order (the player's cleanup running before a Modal
// opened above it) put back a stale "hidden". Invisible today — the app scrolls inside <main>, not
// the document (measured) — but a count cannot be restored out of order, and PlayerModal's
// second copy of the lock is gone.
let scrollLocks = 0;
let unlockedOverflow = "";
export function lockBodyScroll(): () => void {
  if (scrollLocks++ === 0) {
    unlockedOverflow = document.body.style.overflow;
    document.body.style.overflow = "hidden";
  }
  let released = false;
  return () => {
    if (released) return;
    released = true;
    if (--scrollLocks === 0) document.body.style.overflow = unlockedOverflow;
  };
}

/** The rest of the contract: lock the page behind, open focused, hand the focus back on close.
 *  `scopeRef` is the panel the focus should land in; `ownerRef`, when the dialog is bigger than its
 *  panel, is that whole dialog — it must be the SAME element `trapTab` is scoped to, or the two
 *  disagree about what counts as "inside" (see `heldRef`). `focusScope` opens the dialog focused on
 *  the scope element itself rather than its first control: the player keeps the focus on its card, so
 *  Space and the arrows reach its shortcuts instead of activating a button. `initial`, when given,
 *  picks the first-open target instead (`ui/Modal`'s focus-by-purpose rule); null falls back. */
export function useDialogChrome(
  scopeRef: RefObject<HTMLElement | null>,
  ownerRef?: RefObject<HTMLElement | null>,
  {
    focusScope = false,
    initial,
  }: { focusScope?: boolean; initial?: () => HTMLElement | null } = {},
): void {
  // Capture the trigger to restore focus to, during THIS render — i.e. before commit runs the
  // dialog's autoFocus (or our own initial focus) and moves activeElement into the dialog. Reading it
  // later (in the effect) would capture a now-inside element that unmounts on close, stranding focus
  // on <body>. useState initializer runs exactly once, on first render.
  const [restoreTo] = useState(() => document.activeElement as HTMLElement | null);
  // What the dialog is focused ON, so a RE-mount can put the focus back where it was rather than on
  // the first control. React 18's StrictMode runs every effect mount→cleanup→mount in development:
  // the cleanup hands focus back to the trigger, and the second mount then found the focus outside
  // the dialog and moved it to the first focusable — taking it off the `autoFocus` field a beat
  // after the browser had put it there. Measured in Chrome: the command palette opened with the
  // focus on its ✕ in dev and on its search field in prod, from the same code.
  const focusedRef = useRef<HTMLElement | null>(null);
  // Whether the dialog still holds the focus. Read at cleanup rather than computed there: by the
  // time a passive cleanup runs on a real unmount the panel is already out of the document and the
  // focus has fallen to <body>, so "does the panel contain activeElement" answers false for both a
  // close and a user who clicked away — and those two want opposite things.
  //
  // Measured against the DIALOG, not the panel. The two scopes must be the same one the Tab trap
  // uses, and `ui/Modal`'s trap deliberately spans the dialog root so an in-dialog toast's Undo/✕
  // stays reachable — so tabbing to that ✕ read as "the user left the dialog", and since removing
  // the toast fires no `focusin` to re-arm, closing the dialog afterwards stranded the focus on
  // <body> instead of handing it back to the trigger. Measured in jsdom, and a regression: before
  // this hook existed the restore was unconditional.
  const heldRef = useRef(true);
  // Read at open, not a dependency: a caller's fresh function each render must not re-run the open.
  const pickInitial = useEffectEvent(() => initial?.() ?? null);
  useEffect(() => {
    const unlockScroll = lockBodyScroll();
    // Initial focus: unless an autoFocus child already placed it inside (e.g. a rename input), move
    // focus to where this dialog was focused before — or, the first time, to the first control (or
    // the scope itself, with `focusScope`), so the dialog opens focused rather than stranded on the
    // background.
    const scope = scopeRef.current;
    // The dialog's element as it is while open, for the restore in the CLEANUP — where a closing
    // dialog's refs are already detached (refs detach in the commit, passive cleanups run after).
    // Today that reads the same either way (a detached owner and a null one both fail every
    // `contains` below), but the cleanup should not depend on a ref React has already let go of.
    const ownerAtOpen = ownerRef?.current ?? scope;
    if (scope && !scope.contains(document.activeElement)) {
      const back = focusedRef.current;
      const target =
        (back?.isConnected && scope.contains(back) ? back : null) ??
        pickInitial() ??
        (focusScope ? scope : (focusablesIn(scope)[0] ?? scope));
      target.focus();
    }
    const onFocusIn = () => {
      const owner = ownerRef?.current ?? scopeRef.current;
      const active = document.activeElement;
      // "Still ours" also means a dialog stacked ABOVE us. A nested `ui/Modal` portals into its own
      // host on <body>, so it is outside our root BY CONSTRUCTION — widening the scope cannot see
      // it — and while it held the focus we read that as the user walking away. Then both closing
      // together left the focus on <body>: ours skipped the restore, and the child's own
      // `restoreTo` was already detached (passive cleanups run parent-first). Measured in jsdom,
      // and it is the ordinary case — an error dialog raised over a dialog.
      heldRef.current =
        (owner?.contains(active) ?? false) ||
        (active instanceof Element && active.closest('[aria-modal="true"]') != null);
      // Where to put the focus back on a re-mount, though: the PANEL's own last focus. A toast is
      // gone by then, so restoring onto it would strand the focus exactly as before.
      if (scopeRef.current?.contains(document.activeElement))
        focusedRef.current = document.activeElement as HTMLElement;
    };
    onFocusIn();
    document.addEventListener("focusin", onFocusIn);
    return () => {
      document.removeEventListener("focusin", onFocusIn);
      unlockScroll();
      // Restore focus to the trigger, if it's still in the document and focusable — and only if the
      // dialog still held the focus. Once the user has moved it elsewhere, taking it back is the
      // rudest thing a closing dialog can do.
      //
      // Not, though, when a dialog stacked above us is STILL STANDING. `heldRef` counts focus in any
      // `aria-modal` subtree as ours (see above), which cannot tell a layer closing WITH us from one
      // that stays open — and the layer stack cannot tell them apart either: at this instant the
      // child's `removeLayer` has not run, because passive cleanups are parent-first and this hook
      // is declared before that effect, so BOTH cases read "a layer is above me" (measured). What
      // does separate them is the DOM: React detaches a deleted subtree before passive cleanups, so
      // a layer closing with us is already disconnected while one that stays open is not. Without
      // this, a dialog closing on its own (an in-flight mutation's `onSuccess`) under a modal raised
      // above it pulled the focus out of that standing dialog onto our own trigger, behind its
      // scrim — where Enter still activated it.
      const owner = ownerAtOpen;
      const active = document.activeElement;
      const above = active instanceof Element ? active.closest('[aria-modal="true"]') : null;
      const standingAbove =
        above !== null &&
        above.isConnected &&
        !above.contains(owner as Node | null) &&
        !(owner?.contains(above) ?? false);
      if (heldRef.current && !standingAbove && restoreTo && restoreTo.isConnected)
        restoreTo.focus();
    };
  }, [restoreTo, scopeRef, ownerRef, focusScope]);
}