siftlode/frontend/src/lib/useFocusKeeper.ts

import { useEffect, type RefObject } from "react";

// When the element holding the keyboard focus is REMOVED by what it just did — an invite row gone
// after Approve, a toolbar gone after "Clear all", a section gone once its last row is handled — the
// browser drops the focus to <body>, and the next Tab starts over at the top of the document. That is
// the same stranding FocusSafeButton fixes for a button that merely turns disabled, one step further.
//
// So the page keeps the focus in place: the first focusable thing now in the removed node's slot —
// in what followed it, else what preceded it, else the nearest ancestor still on the page. Scoped to
// one container (the page scroller), which is what keeps it off the layers with their own focus
// logic: a menu or a dialog is portaled outside it and returns the focus to its opener itself.
// Keyboard lists (`listNav`) are left alone as well — their `j` already goes on from the slot, and a
// row that merely scrolled out of a virtualized list must not pull the focus after it.

const FOCUSABLE =
  'button:not([disabled]), a[href], input:not([disabled]), select:not([disabled]), textarea:not([disabled]), [tabindex]:not([tabindex="-1"])';

// A focusable element the CSS has taken off screen (a `tabletHidden` column's header button) refuses
// `focus()`, and the focus it was meant to catch fell to <body> (measured in Chrome). An engine
// without `checkVisibility` (jsdom) paints nothing, so there everything counts as shown.
const shown = (el: HTMLElement) => !el.checkVisibility || el.checkVisibility();

function firstFocusable(node: Node | null, fromEnd = false): HTMLElement | null {
  let n: Node | null = node;
  while (n) {
    if (n instanceof HTMLElement) {
      if (n.matches(FOCUSABLE) && shown(n)) return n;
      const all = [...n.querySelectorAll<HTMLElement>(FOCUSABLE)].filter(shown);
      const hit = fromEnd ? all[all.length - 1] : all[0];
      if (hit) return hit;
    }
    n = fromEnd ? n.previousSibling : n.nextSibling;
  }
  return null;
}

/** Move the focus off `el`, which is about to go, to its nearest focusable neighbour — the next one,
 *  else the previous one, widening to the parent's neighbours — so its removal does not drop the
 *  focus to <body>. For a removal the caller makes itself (a chip's ×) anywhere, including a modal
 *  or a side panel, where the keeper below does not watch. False when nothing was found. */
export function focusNeighbour(el: HTMLElement): boolean {
  for (let at: HTMLElement | null = el; at && at !== document.body; at = at.parentElement) {
    const hit = firstFocusable(at.nextSibling) ?? firstFocusable(at.previousSibling, true);
    if (hit) {
      hit.focus({ preventScroll: true });
      return true;
    }
  }
  return false;
}

/** Run `remove`, which takes `el` off the page, without dropping `el`'s focus: it goes to the
 *  neighbour first (`focusNeighbour`). That neighbour may go in the same removal — the "Clear all"
 *  that only shows beside two filters, the "+1" of the last hidden one — so
 *  once the removal has rendered, a focus left on <body> goes to what preceded `el`, else to the
 *  first control left in its row. Only when the focus was lost: a focus the reader moved stays. */
export function handOffFocus(el: HTMLElement, remove: () => void): void {
  const before = firstFocusable(el.previousSibling, true);
  const row = el.parentElement;
  focusNeighbour(el);
  remove();
  // A timer, not a frame: a page the browser is not painting runs no requestAnimationFrame.
  window.setTimeout(() => {
    const now = document.activeElement;
    if (now && now !== document.body && now.isConnected) return;
    const stayed = before?.isConnected && shown(before) ? before : null;
    (stayed ?? (row?.isConnected ? firstFocusable(row.firstChild) : null))?.focus({
      preventScroll: true,
    });
  }, 0);
}

// The element the keeper last put the focus on, while it still holds it. A Back that swaps the page
// or a sub-view removes the held control too, and the keeper's stand-in then sat where
// `restoreNavPosition` looks for an untouched <body>: the restore gave up and Back left the focus on
// the first tile instead of the one the view was opened from (measured in Chrome). A stand-in is
// not the reader's choice, so the restore may take the focus from it.
let standIn: HTMLElement | null = null;

/** Whether `el` holds the focus only because the keeper put it there. */
function isStandInFocus(el: Element | null): boolean {
  return el !== null && el === standIn;
}

/** Marks an element that can hold the focus without being a control: the page's scroll region
 *  (`<main tabIndex={-1}>`, the skip link's target) takes it from a mouse click on the page's blank
 *  area. Holding it there is the focus RESTING on the page, the same as <body>. */
export const FOCUS_REST_ATTRS = { "data-focus-rest": "" } as const;

/** Whether the focus is nowhere the reader chose — <body>, the keeper's stand-in, or a resting place
 *  (`FOCUS_REST_ATTRS`) — so a restore may take it. The page region used to count as a choice: a
 *  click on its background focused it, and Back then left the focus there instead of on the item
 *  the view was opened from. */
export function focusIsUnclaimed(el: Element | null = document.activeElement): boolean {
  return !el || el === document.body || isStandInFocus(el) || el.hasAttribute("data-focus-rest");
}

/** Where the focus goes once `lost` (removed with the nodes of `record`) has left the page — never
 *  outside `root`: a slot with nothing focusable left (a page swapped for its loading fallback) must
 *  not send the focus to the first control of the whole document. */
function slotFor(record: MutationRecord, lost: Node, root: HTMLElement): HTMLElement | null {
  const removed = [...record.removedNodes].find((n) => n === lost || n.contains(lost));
  if (!removed) return null;
  const after = firstFocusable(record.nextSibling);
  if (after) return after;
  const before = firstFocusable(record.previousSibling, true);
  if (before) return before;
  let up: Node | null = record.target;
  while (up && up instanceof HTMLElement && root.contains(up)) {
    const inside = firstFocusable(up.firstChild);
    if (inside) return inside;
    up = up.parentElement;
  }
  return null;
}

export function useFocusKeeper(ref: RefObject<HTMLElement | null>): void {
  useEffect(() => {
    const root = ref.current;
    if (!root) return;
    let held: HTMLElement | null = null;
    const onFocusIn = (e: FocusEvent) => {
      const t = e.target as HTMLElement | null;
      if (t !== standIn) standIn = null;
      held = t && root.contains(t) && !t.closest("[data-nav-index]") ? t : null;
    };
    // Using the stand-in makes it the reader's own focus: a click on it or a key typed into it fires
    // no focusin, and a restore still polling for its list would otherwise take the focus from
    // under the reader's hands.
    const onUse = () => {
      standIn = null;
    };
    const watch = new MutationObserver((records) => {
      const lost = held;
      if (!lost || lost.isConnected) return;
      held = null;
      if (!focusIsUnclaimed()) return; // something already took the focus
      for (const r of records) {
        const target = slotFor(r, lost, root);
        if (target) {
          standIn = target;
          target.focus({ preventScroll: true });
          return;
        }
      }
    });
    document.addEventListener("focusin", onFocusIn);
    document.addEventListener("pointerdown", onUse, true);
    document.addEventListener("keydown", onUse, true);
    watch.observe(root, { childList: true, subtree: true });
    return () => {
      document.removeEventListener("focusin", onFocusIn);
      document.removeEventListener("pointerdown", onUse, true);
      document.removeEventListener("keydown", onUse, true);
      watch.disconnect();
    };
  }, [ref]);
}