import { type KeyboardEvent, useEffect, useState } from "react";

const NEXT_KEYS = ["ArrowDown", "ArrowRight"];
const PREV_KEYS = ["ArrowUp", "ArrowLeft"];
const HORIZONTAL_KEYS = ["ArrowLeft", "ArrowRight"];
const ARROW_KEYS = [...HORIZONTAL_KEYS, "ArrowUp", "ArrowDown"];

/**
 * Find the item closest to `from` in the neighbouring visual row or column.
 *
 * The app list is a grid whose column count depends on the available width, so
 * which row and column an item ends up in is only known after layout. Comparing
 * client rectangles is therefore the only reliable way to move in two
 * dimensions.
 */
function findAdjacent(items: readonly HTMLElement[], from: number, axis: "x" | "y", direction: 1 | -1) {
	// `main` is the axis stepped along, `cross` the one used to pick the closest
	// candidate within the target row/column.
	let rects = items.map((item) => {
		let rect = item.getBoundingClientRect();
		return axis === "x" ? { main: rect.left, cross: rect.top } : { main: rect.top, cross: rect.left };
	});

	let origin = rects[from];
	if (!origin) return null;

	// The nearest row/column edge in the requested direction.
	let target: number | null = null;
	for (let rect of rects) {
		let delta = rect.main - origin.main;
		if (direction > 0 ? delta <= 1 : delta >= -1) continue;
		if (target === null || Math.abs(rect.main - origin.main) < Math.abs(target - origin.main)) {
			target = rect.main;
		}
	}

	if (target === null) return null;

	let best: number | null = null;
	let bestDistance = Number.POSITIVE_INFINITY;
	for (let index = 0; index < rects.length; index++) {
		let rect = rects[index];
		if (Math.abs(rect.main - target) > 1) continue;

		let distance = Math.abs(rect.cross - origin.cross);
		if (distance < bestDistance) {
			bestDistance = distance;
			best = index;
		}
	}

	return best;
}

export type RovingTabIndexOptions = {
	/**
	 * When true, the arrow keys follow the visual layout of a wrapping grid
	 * instead of the DOM order: left/right move between columns and up/down
	 * between rows.
	 */
	columns?: boolean;
};

/**
 * Roving tabindex: the container holds a single tab stop, and the arrow keys
 * (plus Home/End) move focus between the items inside it.
 *
 * Spread `getItemProps(index)` on every focusable item and put `onKeyDown` on
 * their common ancestor.
 */
export function useRovingTabIndex(itemCount: number, { columns = false }: RovingTabIndexOptions = {}) {
	let [activeIndex, setActiveIndex] = useState(0);

	// Keep the tab stop within the list when it shrinks (e.g. while searching).
	useEffect(() => {
		setActiveIndex((index) => (index >= itemCount ? 0 : index));
	}, [itemCount]);

	function onKeyDown(event: KeyboardEvent<HTMLElement>) {
		if (itemCount === 0) return;

		let items = [...event.currentTarget.querySelectorAll<HTMLElement>("[data-roving-item]")];
		let next: number | null;

		if (columns && ARROW_KEYS.includes(event.key)) {
			event.preventDefault();

			let isHorizontal = HORIZONTAL_KEYS.includes(event.key);
			let direction: 1 | -1 = NEXT_KEYS.includes(event.key) ? 1 : -1;
			next = findAdjacent(items, activeIndex, isHorizontal ? "x" : "y", direction);

			if (next === null) {
				// Stay put at the first/last column, but let up/down continue into the
				// neighbouring item so the whole list stays reachable.
				if (isHorizontal) return;
				next = (activeIndex + direction + itemCount) % itemCount;
			}
		} else if (NEXT_KEYS.includes(event.key)) next = (activeIndex + 1) % itemCount;
		else if (PREV_KEYS.includes(event.key)) next = (activeIndex - 1 + itemCount) % itemCount;
		else if (event.key === "Home") next = 0;
		else if (event.key === "End") next = itemCount - 1;
		else return;

		event.preventDefault();
		setActiveIndex(next);
		items[next]?.focus();
	}

	function getItemProps(index: number) {
		return {
			"data-roving-item": "",
			tabIndex: index === activeIndex ? 0 : -1,
			onFocus: () => setActiveIndex(index),
		};
	}

	/**
	 * Props for a secondary control belonging to item `index`, e.g. the favorite
	 * toggle next to an app link. It shares the item's tab stop — so Tab reaches
	 * it from the link — but the arrow keys skip it.
	 */
	function getSecondaryProps(index: number) {
		return {
			tabIndex: index === activeIndex ? 0 : -1,
			onFocus: () => setActiveIndex(index),
		};
	}

	return { onKeyDown, getItemProps, getSecondaryProps };
}