Key hints

"use client";

import {
  cloneElement,
  createContext,
  isValidElement,
  useContext,
  useEffect,
  useId,
  useMemo,
  useRef,
  type ComponentProps,
  type PointerEvent,
  type ReactElement,
  type ReactNode,
  type Ref,
} from "react";
import "./key-hints.css";

type Modifier = "Alt" | "Control" | "Meta";
type Entry = { keys: string; target: () => HTMLElement | null };
type Registry = {
  modifier: Modifier;
  register: (id: string, entry: Entry) => () => void;
};

const KeyHintsContext = createContext<Registry | null>(null);

/** True when the chosen modifier is the only one held, so Alt+Shift (or Ctrl+Alt) is left alone. */
function held(event: KeyboardEvent, modifier: Modifier) {
  const down = { Alt: event.altKey, Control: event.ctrlKey, Meta: event.metaKey };
  return down[modifier] && !event.shiftKey && (Object.keys(down) as Modifier[]).every((m) => m === modifier || !down[m]);
}

/**
 * Single characters match on the physical key (event.code), because a held Option on macOS
 * turns "s" into "ß". Named keys such as Enter match on event.key.
 */
function matches(event: KeyboardEvent, keys: string) {
  if (keys.length === 1) {
    const ch = keys.toUpperCase();
    if (/[A-Z]/.test(ch)) return event.code === `Key${ch}`;
    if (/[0-9]/.test(ch)) return event.code === `Digit${ch}` || event.code === `Numpad${ch}`;
  }
  return event.key.toLowerCase() === keys.toLowerCase();
}

function label(keys: string) {
  return keys.length === 1 ? keys.toUpperCase() : keys;
}

function assignRef<T>(ref: Ref<T> | undefined, node: T | null) {
  if (typeof ref === "function") ref(node);
  else if (ref) ref.current = node;
}

function usable(el: HTMLElement) {
  return !el.matches(":disabled") && el.getAttribute("aria-disabled") !== "true";
}

export type KeyHintsProps = ComponentProps<"div"> & {
  /** The key the reader holds to show the hints and use them. */
  modifier?: Modifier;
  children: ReactNode;
};

export function KeyHints({ modifier = "Alt", className, children, ref, onPointerEnter, ...rest }: KeyHintsProps) {
  const root = useRef<HTMLDivElement | null>(null);
  const entries = useRef(new Map<string, Entry>());
  // One registry per instance; hints register into this instance's own Map.
  const registry = useMemo<Registry>(
    () => ({
      modifier,
      register: (id, entry) => {
        entries.current.set(id, entry);
        return () => {
          entries.current.delete(id);
        };
      },
    }),
    [modifier],
  );

  useEffect(() => {
    const own = root.current;
    if (!own) return;

    const show = (on: boolean) => own.toggleAttribute("data-active", on);

    /*
     * Every KeyHints with this modifier shows its caps. Only one acts on a key: the one holding
     * focus, else the one the pointer entered last, else the first in the document.
     * Each instance works this out from the DOM, so nothing is shared between them.
     */
    const winner = (event: KeyboardEvent) => {
      const candidates = [...document.querySelectorAll<HTMLElement>(".eq-key-hints")].filter(
        (r) =>
          r.dataset.modifier === modifier &&
          [...r.querySelectorAll<HTMLElement>(".eq-key-hints-target")].some((t) => {
            const control = t.firstElementChild;
            return (
              t.dataset.keys != null &&
              matches(event, t.dataset.keys) &&
              control instanceof HTMLElement &&
              usable(control)
            );
          }),
      );
      const focus = document.activeElement;
      const focused = candidates.filter((r) => focus && r.contains(focus));
      if (focused.length) return focused[focused.length - 1];
      let best: HTMLElement | undefined;
      for (const r of candidates) {
        const stamp = Number(r.dataset.pointerAt ?? -1);
        if (!best || stamp > Number(best.dataset.pointerAt ?? -1)) best = r;
      }
      return best;
    };

    const onKeyDown = (event: KeyboardEvent) => {
      const on = held(event, modifier);
      show(on);
      if (!on || event.repeat || event.key === modifier) return;
      const hit = [...entries.current.values()].find((e) => matches(event, e.keys));
      const target = hit?.target();
      if (!target || !usable(target) || winner(event) !== own) return;
      // Capturing: stop the page's own single-key shortcuts from also seeing this key.
      event.preventDefault();
      event.stopPropagation();
      target.focus();
      target.click();
    };
    const onKeyUp = (event: KeyboardEvent) => show(held(event, modifier));
    const onBlur = () => show(false);

    window.addEventListener("keydown", onKeyDown, true);
    window.addEventListener("keyup", onKeyUp, true);
    window.addEventListener("blur", onBlur);
    return () => {
      window.removeEventListener("keydown", onKeyDown, true);
      window.removeEventListener("keyup", onKeyUp, true);
      window.removeEventListener("blur", onBlur);
      own.removeAttribute("data-active");
    };
  }, [modifier]);

  return (
    <KeyHintsContext.Provider value={registry}>
      <div
        {...rest}
        ref={(node) => {
          root.current = node;
          assignRef(ref, node);
        }}
        className={["eq-key-hints", className].filter(Boolean).join(" ")}
        data-modifier={modifier}
        onPointerEnter={(event: PointerEvent<HTMLDivElement>) => {
          onPointerEnter?.(event);
          event.currentTarget.dataset.pointerAt = String(event.timeStamp);
        }}
      >
        {children}
      </div>
    </KeyHintsContext.Provider>
  );
}

export type KeyHintProps = ComponentProps<"span"> & {
  /** One character such as "s", or a key name such as "Enter". */
  keys: string;
  /** Exactly one focusable element, such as a button or a link. */
  children: ReactNode;
};

export function KeyHint({ keys, className, children, ref, ...rest }: KeyHintProps) {
  const id = useId();
  const wrap = useRef<HTMLSpanElement | null>(null);
  const context = useContext(KeyHintsContext);
  const modifier = context?.modifier ?? "Alt";
  const shortcut = `${modifier}+${label(keys)}`;

  useEffect(() => {
    if (!context) return;
    return context.register(id, {
      keys,
      target: () => {
        const el = wrap.current?.firstElementChild;
        return el instanceof HTMLElement && !el.classList.contains("eq-key-hints-cap") ? el : null;
      },
    });
  }, [context, id, keys]);

  const child = isValidElement<{ "aria-keyshortcuts"?: string }>(children)
    ? cloneElement(children as ReactElement<{ "aria-keyshortcuts"?: string }>, { "aria-keyshortcuts": shortcut })
    : children;

  return (
    <span
      {...rest}
      ref={(node) => {
        wrap.current = node;
        assignRef(ref, node);
      }}
      className={["eq-key-hints-target", className].filter(Boolean).join(" ")}
      data-keys={keys}
      aria-keyshortcuts={child === children ? shortcut : undefined}
    >
      {child}
      <kbd className="eq-key-hints-cap" aria-hidden="true">
        {label(keys)}
      </kbd>
    </span>
  );
}