import {

  findAll,

  findAllFocusable,

  findFocusableByIndexOrLast,

  findAllVisible,

  isElement,

  isElementFocusable,

  Lazy,

type GridDirection = "up" | "down" | "left" | "right";

type SequenceDirection = "prev" | "next";

class GroupModel {

  readonly #groupElement;

  #cursorElement?: HTMLElement | SVGElement;

  #cursor?: number | [number, number];

interface GroupModel {

  readonly type: string;

  readonly container: HTMLElement;

  constructor(groupElement: HTMLElement | SVGElement) {

    this.#groupElement = groupElement;

  /** Returns true if it was possible to set the given element as current. */

  setCurrent(focusableElement: HTMLElement): boolean;

  /** Returns the element that was focused, if any. */

  focusCurrent(): HTMLElement | undefined;

  /** Returns true if the focus was moved. */

  moveFocus(direction: GridDirection): boolean;

}

  get #groupType() {

    const type = this.#groupElement.dataset.kbdGroup;

    if (type === "vertical" || type === "horizontal" || type === "grid") {

      return type;

const SEL_FOCUSABLE = "[data-kbd-focusable]";

const SEL_GRID_ROW = "[data-kbd-grid-row]";

const SEL_GROUP = "[data-kbd-group]";

const DATA_GROUP_TYPE = "kbdGroup";

const GROUP_TYPE_GRID = "grid";

const GROUP_TYPE_HORIZONTAL = "horizontal";

const GROUP_TYPE_VERTICAL = "vertical";

interface GridCurrent {

  cell: HTMLElement;

  row: HTMLElement;

  x: number;

  y: number;

  requestedX?: number;

  rows: HTMLElement[];

  cells: HTMLElement[];

}

class GridModel implements GroupModel {

  #current?: GridCurrent;

  readonly type = GROUP_TYPE_GRID;

    return undefined;

  constructor(readonly container: HTMLElement) {

  }

  setCursor(focusableElement: EventTarget | null) {

    // Already handled, nothing to do.

    if (this.#cursorElement === focusableElement) return;

  setCurrent(focusableElement: HTMLElement): boolean {

    if (!isElementFocusable(focusableElement)) return false;

    if (!isElement(focusableElement)) return;

    if (!focusableElement.matches(SEL_FOCUSABLE)) return false;

    if (!focusableElement.matches("[data-kbd-focusable]")) return;

    const container = this.container;

    if (!container.contains(focusableElement)) return false;

    const groupElement = this.#groupElement;

    if (!groupElement.contains(focusableElement)) return;

    const row = focusableElement.closest(SEL_GRID_ROW);

    if (!isElement(row)) return false;

    const groupType = this.#groupType;

    if (groupType === "vertical" || groupType === "horizontal") {

      const focusables = findAllFocusable(groupElement, "[data-kbd-focusable]");

      const index = focusables.indexOf(focusableElement);

      if (index === -1) return;

    const rows = findAllVisible(container, SEL_GRID_ROW);

    const cells = findAllFocusable(row, SEL_FOCUSABLE);

      this.#cursor = index;

    } else if (groupType === "grid") {

      const row = focusableElement.closest("[data-kbd-grid-row]");

      if (!isElement(row)) return;

    this.#current = {

      cell: focusableElement,

      row: row,

      rows,

      cells,

      x: cells.indexOf(focusableElement),

      y: rows.indexOf(row),

      requestedX: this.#current?.cell === focusableElement

        ? this.#current.requestedX

        : undefined,

    };

    return true;

  }

      const rows = findAll(groupElement, "[data-kbd-grid-row]");

      const rowFocusables = findAllFocusable(row, "[data-kbd-focusable]");

      const rowIndex = rows.indexOf(row);

      const colIndex = rowFocusables.indexOf(focusableElement);

  #getCurrent(): GridCurrent | undefined {

    const container = this.container;

    let { cell, row, x, y, requestedX } = this.#current ??

      { x: 0, y: 0 };

      if (rowIndex === -1 || colIndex === -1) return;

    const rows = findAllVisible(container, SEL_GRID_ROW);

      this.#cursor = [colIndex, rowIndex];

    if (!row || !rows.includes(row)) {

      row = rows[y] ?? rows.at(-1);

    }

    this.#cursorElement = focusableElement;

    if (!row) return;

    const cells = findAllFocusable(row, SEL_FOCUSABLE);

    if (!cell || !cells.includes(cell)) {

      cell = cells[requestedX ?? x] ?? rows.at(-1);

    }

  #focus(

    focusable: Element | undefined,

    cursor: number | [number, number],

  ): boolean {

    if (!isElement(focusable)) return false;

    if (!cell) return;

    focusable.focus();

    this.#cursor = cursor;

    this.#cursorElement = focusable;

    return true;

    return {

      cell,

      row,

      rows,

      cells,

      x: cells.indexOf(cell),

      y: rows.indexOf(row),

      requestedX,

    };

  }

  #attemptToFocusSequence(direction: -1 | 1): boolean {

    const groupType = this.#groupType;

    const cursor = this.#cursor;

    if (groupType !== "vertical" && groupType !== "horizontal") return false;

    if (typeof cursor !== "number") return false;

  focusCurrent(): HTMLElement | undefined {

    this.#current = this.#getCurrent();

    this.#current?.cell.focus();

    return this.#current?.cell;

  }

    const focusables = findAllFocusable(

      this.#groupElement,

      "[data-kbd-focusable]",

    );

  #tryMoveFocusUpDown(direction: -1 | 1): boolean {

    const current = this.#getCurrent();

    if (!current) return false;

    const row = current.rows[current.y + direction];

    if (!row) return false;

    return this.#focus(focusables[cursor + direction], cursor + direction);

    const requestedX = current.requestedX ?? current.x;

    const cells = findAllFocusable(row, SEL_FOCUSABLE);

    const cell = cells[requestedX] ?? cells.at(-1);

    if (!cell) return false;

    current.y += direction;

    current.row = row;

    current.cells = cells;

    current.cell = cell;

    current.x = cells.indexOf(cell);

    current.requestedX = requestedX;

    this.#current = current;

    cell.focus();

    return true;

  }

  #attemptToFocusUpDown(direction: -1 | 1): boolean {

    const cursor = this.#cursor;

    if (this.#groupType !== "grid") return false;

    if (!Array.isArray(cursor)) return false;

  #tryMoveFocusLeftRight(direction: -1 | 1): boolean {

    const current = this.#getCurrent();

    if (!current) return false;

    const [cursorX, cursorY] = cursor;

    const cell = current.cells[current.x + direction];

    if (!cell) return false;

    const rows = findAll(this.#groupElement, "[data-kbd-grid-row]");

    const row = rows[cursorY + direction];

    const focusable = findFocusableByIndexOrLast(

      row,

      "[data-kbd-focusable]",

      cursorX,

    );

    return this.#focus(focusable, [cursorX, cursorY + direction]);

    current.cell = cell;

    current.x = current.cells.indexOf(cell);

    current.requestedX = undefined;

    this.#current = current;

    cell.focus();

    return true;

  }

  #attemptToFocusLeftRight(direction: -1 | 1): boolean {

    const cursor = this.#cursor;

    if (this.#groupType !== "grid") return false;

    if (!this.#cursorElement) return false;

    if (!Array.isArray(cursor)) return false;

  moveFocus(direction: GridDirection): boolean {

    if (direction === "up") {

      return this.#tryMoveFocusUpDown(-1);

    } else if (direction === "down") {

      return this.#tryMoveFocusUpDown(1);

    } else if (direction === "left") {

      return this.#tryMoveFocusLeftRight(-1);

    } else if (direction === "right") {

      return this.#tryMoveFocusLeftRight(1);

    }

    const [, cursorY] = cursor;

    return false;

  }

}

    const rows = findAll(this.#groupElement, "[data-kbd-grid-row]");

    const row = rows[cursorY];

    const focusables = findAllFocusable(row, "[data-kbd-focusable]");

    const cursorX = focusables.indexOf(this.#cursorElement);

    const focusable = focusables[cursorX + direction];

    return this.#focus(focusable, [cursorX + direction, cursorY]);

interface SequenceCurrent {

  cell: HTMLElement;

  cells: HTMLElement[];

  index: number;

}

class SequenceModel implements GroupModel {

  #current?: SequenceCurrent;

  constructor(

    readonly type: typeof GROUP_TYPE_HORIZONTAL | typeof GROUP_TYPE_VERTICAL,

    readonly container: HTMLElement,

  ) {}

  focusCurrent(): boolean {

    if (this.#cursorElement) {

      this.#cursorElement.focus();

      return true;

    } else {

      const firstElement = this.#groupElement.querySelector(

        "[data-kbd-focusable]",

      );

  setCurrent(focusableElement: HTMLElement): boolean {

    if (!isElementFocusable(focusableElement)) return false;

    if (!focusableElement.matches(SEL_FOCUSABLE)) return false;

      if (isElement(firstElement)) {

        this.setCursor(firstElement);

        firstElement.focus();

    const container = this.container;

    if (!container.contains(focusableElement)) return false;

    const cells = findAllFocusable(container, SEL_FOCUSABLE);

    this.#current = {

      cell: focusableElement,

      cells,

      index: cells.indexOf(focusableElement),

    };

    return true;

  }

  #getCurrent(): SequenceCurrent | undefined {

    const container = this.container;

    let { cell, index } = this.#current ?? { index: 0 };

    const cells = findAllFocusable(container, SEL_FOCUSABLE);

    if (!cell || !cells.includes(cell)) {

      cell = cells[index] ?? cells.at(-1);

    }

    if (!cell) return;

    return { cell, cells, index: cells.indexOf(cell) };

  }

    return false;

  focusCurrent(): HTMLElement | undefined {

    this.#current = this.#getCurrent();

    this.#current?.cell.focus();

    return this.#current?.cell;

  }

  moveFocus(direction: GridDirection): boolean {

    const groupType = this.#groupType;

  #tryMoveFocus(direction: -1 | 1): boolean {

    const current = this.#getCurrent();

    if (!current) return false;

    const cell = current.cells[current.index + direction];

    if (!cell) return false;

    current.cell = cell;

    current.index = current.cells.indexOf(cell);

    this.#current = current;

    cell.focus();

    return true;

  }

  moveFocus(direction: GridDirection): boolean {

    const type = this.type;

    if (

      (groupType === "vertical" && direction === "up") ||

      (groupType === "horizontal" && direction === "left")

      (direction === "left" && type == GROUP_TYPE_HORIZONTAL) ||

      (direction === "up" && type === GROUP_TYPE_VERTICAL)

    ) {

      return this.#attemptToFocusSequence(-1);

      return this.#tryMoveFocus(-1);

    } else if (

      (groupType === "vertical" && direction === "down") ||

      (groupType === "horizontal" && direction === "right")

      (direction === "right" && type == GROUP_TYPE_HORIZONTAL) ||

      (direction === "down" && type === GROUP_TYPE_VERTICAL)

    ) {

      return this.#attemptToFocusSequence(1);

    } else if (groupType === "grid") {

      if (direction === "up") {

        return this.#attemptToFocusUpDown(-1);

      } else if (direction === "down") {

        return this.#attemptToFocusUpDown(1);

      } else if (direction === "left") {

        return this.#attemptToFocusLeftRight(-1);

      } else if (direction === "right") {

        return this.#attemptToFocusLeftRight(1);

      }

      return this.#tryMoveFocus(1);

    }

    return false;

}

class DocumentModel {

  readonly #groupModels = new WeakMap<Element, GroupModel>();

  readonly #groups = new WeakMap<Element, GroupModel>();

  readonly #creators: Record<

    string,

    ((container: HTMLElement) => GroupModel) | undefined

  > = {

    [GROUP_TYPE_GRID]: (x) => new GridModel(x),

    [GROUP_TYPE_HORIZONTAL]: (x) => new SequenceModel(GROUP_TYPE_HORIZONTAL, x),

    [GROUP_TYPE_VERTICAL]: (x) => new SequenceModel(GROUP_TYPE_VERTICAL, x),

  };

  getGroupModel(element: EventTarget | null): GroupModel | undefined {

    if (!isElement(element)) return;

    const groupElement = element.closest("[data-kbd-group]");

    const groupElement = element.closest(SEL_GROUP);

    if (!isElement(groupElement)) return;

    let groupData = this.#groupModels.get(groupElement);

    if (!groupData) {

      groupData = new GroupModel(groupElement);

      this.#groupModels.set(groupElement, groupData);

    const groupType = groupElement.dataset[DATA_GROUP_TYPE];

    let groupModel = this.#groups.get(groupElement);

    if (!groupModel || groupModel.type !== groupType) {

      groupModel = groupType

        ? this.#creators[groupType]?.(groupElement)

        : undefined;

      if (groupModel) {

        this.#groups.set(groupElement, groupModel);

      } else {

        this.#groups.delete(groupElement);

      }

    }

    return groupData;

    return groupModel;

  }

  getAdjacentGroupModel(

  ): GroupModel | undefined {

    if (!isElement(element)) return;

    const groupElement = element.closest("[data-kbd-group]");

    const groupElement = element.closest(SEL_GROUP);

    if (!isElement(groupElement)) return;

    const groupElements = findAll(document.body, "[data-kbd-group]");

    const groupElements = findAll(document.body, SEL_GROUP);

    const length = groupElements.length;

    const index = groupElements.indexOf(groupElement);

    const dir = direction === "prev" ? -1 : 1;

  readonly #subTreeObserver = new Lazy(() =>

    new MutationObserver(() => {

      if (!this.#lastFocused) return;

      const { element, getReplacement, gridTargetX } = this.#lastFocused;

      if (isElementFocusable(element)) return;

      const replacement = getReplacement();

      if (!replacement || !document.contains(replacement)) return;

      this.#focusElement({ element: replacement, getReplacement, gridTargetX });

      if (this.#lastFocused && !isElementFocusable(this.#lastFocused.element)) {

        const newElement = this.#lastFocused.group.focusCurrent();

        if (newElement) {

          this.#lastFocused = {

            group: this.#lastFocused.group,

            element: newElement,

          };

        }

      }

    })

  );

  #lastFocused?: {

    element: Element;

    getReplacement: () => Element | undefined;

    gridTargetX?: number;

    group: GroupModel;

    element: HTMLElement;

  };

  #focusElement(

    { event, element, getReplacement, gridTargetX }: {

      event?: KeyboardEvent;

      element: Element | undefined;

      getReplacement: () => Element | undefined;

      gridTargetX: number | undefined;

    },

  ) {

    if (!isElement(element)) return;

    event?.preventDefault();

    this.#lastFocused = { element, getReplacement, gridTargetX };

    element.focus();

  #focusListener = (e: FocusEvent) => {

    this.#subTreeObserver.value.disconnect();

    this.#subTreeObserver.value.observe(document, {

    const element = e.target;

    if (element instanceof HTMLElement) {

      const group = this.#documentModel.getGroupModel(element);

      if (group) {

        const didSet = group.setCurrent(element);

        if (didSet) {

          this.#lastFocused = { group, element };

          this.#subTreeObserver.value.observe(group.container, {

            childList: true,

            subtree: true,

            attributes: true,

          });

    element.addEventListener(

      "blur",

      () => {

        if (this.#lastFocused?.element === element) {

          this.#lastFocused = undefined;

          this.#subTreeObserver.value.disconnect();

        }

      },

      { once: true },

    );

      }

  #focusListener = (e: FocusEvent) => {

    this.#documentModel.getGroupModel(e.target)?.setCursor(e.target);

    }

  };

  static #keyDirectionMap: Record<string, GridDirection> = {

    let didFocus: boolean | undefined;

    if (key === "Tab") {

      didFocus = this.#documentModel.getAdjacentGroupModel(

      didFocus = !!this.#documentModel.getAdjacentGroupModel(

        e.target,

        e.shiftKey ? "prev" : "next",

      )

export function isElement(x: unknown): x is HTMLElement | SVGElement {

  return x instanceof HTMLElement || x instanceof SVGElement;

export function isElement(x: unknown): x is HTMLElement {

  return x instanceof HTMLElement;

}

export function findAll(container: Element, selector: string): Element[] {

  return [...container.querySelectorAll(selector)];

export function findAll(

  container: HTMLElement,

  selector: string,

): HTMLElement[] {

  return [...container.querySelectorAll(selector)] as HTMLElement[];

}

export function findAllVisible(

  container: HTMLElement,

  selector: string,

): HTMLElement[] {

  return findAll(container, selector).filter(isElementVisible);

}

export function findAllFocusable(

  container: Element,

  container: HTMLElement,

  selector: string,

): Element[] {

): HTMLElement[] {

  return findAll(container, selector).filter(isElementFocusable);

}

export function findFocusableByIndexOrLast(

  container: Element | undefined,

  selector: string,

  index: number,

): Element | undefined {

  if (!container) return;

  const elements = findAllFocusable(container, selector);

  return elements[index] ?? elements[elements.length - 1];

export function isElementVisible(element: HTMLElement): boolean {

  if (!element.offsetParent) return false;

  const style = getComputedStyle(element);

  if (style.display === "none" || style.visibility === "hidden") return false;

  return true;

}

/**

 * Determine whether an element can be focused.

 */

export function isElementFocusable(element: Element): boolean {

  if (!isElement(element)) return false;

  if (!element.offsetParent) return false;

export function isElementFocusable(element: HTMLElement): boolean {

  if (!isElementVisible(element)) return false;

  if (

    element instanceof HTMLButtonElement ||

    if (element.tabIndex === -1 && !element.isContentEditable) return false;

  }

  // Check if the element is not hidden by CSS.

  const style = getComputedStyle(element);

  if (style.display === "none" || style.visibility === "hidden") return false;

  return true;

}