Selection area

"use client";

import {
  createContext,
  useContext,
  useEffect,
  useRef,
  useState,
  type ComponentProps,
  type KeyboardEvent,
  type MouseEvent,
  type PointerEvent,
} from "react";
import "./selection-area.css";

export type SelectionAreaProps = ComponentProps<"div"> & {
  /** Names the set of items for assistive tech. */
  label: string;
  /** Selected item ids, when you control them. */
  selected?: string[];
  /** Selected item ids on first render, when you do not. */
  defaultSelected?: string[];
  onSelectedChange?: (selected: string[]) => void;
};

export type SelectionItemProps = Omit<ComponentProps<"div">, "id"> & {
  /** Identifies the item in the selection. Not used as the DOM id. */
  id: string;
  disabled?: boolean;
};

type Area = {
  selected: ReadonlySet<string>;
  focusId: string | null;
  pick: (id: string, event: MouseEvent) => void;
  focus: (id: string) => void;
};

const AreaContext = createContext<Area | null>(null);

/* Distance a press must travel before it becomes a marquee. */
const THRESHOLD = 4;
/* Band along each edge where a drag scrolls the area, and the fastest it goes per frame. */
const EDGE = 28;
const SPEED = 14;

type Box = { x: number; y: number; w: number; h: number };
type Drag = {
  pointerId: number;
  x: number;
  y: number;
  clientX: number;
  clientY: number;
  base: string[];
  mode: "replace" | "add" | "toggle";
  moved: boolean;
  frame: number;
};

const same = (a: readonly string[], b: readonly string[]) =>
  a.length === b.length && a.every((id) => b.includes(id));

export function SelectionArea({
  label,
  selected: controlled,
  defaultSelected = [],
  onSelectedChange,
  className,
  children,
  onPointerDown,
  onPointerMove,
  onPointerUp,
  onPointerCancel,
  onKeyDown,
  ...rest
}: SelectionAreaProps) {
  const [uncontrolled, setUncontrolled] = useState(defaultSelected);
  const selected = controlled ?? uncontrolled;
  const [focusId, setFocusId] = useState<string | null>(null);
  const [marquee, setMarquee] = useState<Box | null>(null);
  const rootRef = useRef<HTMLDivElement>(null);
  const drag = useRef<Drag | null>(null);
  const anchor = useRef<string | null>(null);
  const latest = useRef(selected);
  latest.current = selected;

  useEffect(() => () => cancelAnimationFrame(drag.current?.frame ?? 0), []);

  /* Keep exactly one item in the tab order: the focused one, else the first selected, else the first. */
  useEffect(() => {
    const els = items();
    if (focusId !== null && els.some((el) => el.dataset.id === focusId)) return;
    const first = els.find((el) => selected.includes(el.dataset.id ?? "")) ?? els[0];
    setFocusId(first?.dataset.id ?? null);
  });

  function items() {
    const root = rootRef.current;
    if (!root) return [];
    return [...root.querySelectorAll<HTMLElement>(".eq-selection-area-item")].filter(
      (el) => el.closest(".eq-selection-area") === root,
    );
  }

  const enabledIds = () =>
    items()
      .filter((el) => el.getAttribute("aria-disabled") !== "true")
      .map((el) => el.dataset.id ?? "");

  function commit(next: string[]) {
    if (same(next, latest.current)) return;
    latest.current = next;
    if (controlled === undefined) setUncontrolled(next);
    onSelectedChange?.(next);
  }

  function moveFocus(el: HTMLElement | undefined) {
    if (!el?.dataset.id) return;
    setFocusId(el.dataset.id);
    el.focus();
    el.scrollIntoView({ block: "nearest", inline: "nearest" });
  }

  function pick(id: string, event: MouseEvent) {
    const current = latest.current;
    if (event.metaKey || event.ctrlKey) {
      commit(current.includes(id) ? current.filter((x) => x !== id) : [...current, id]);
      anchor.current = id;
    } else if (event.shiftKey && anchor.current !== null) {
      const ids = enabledIds();
      const a = ids.indexOf(anchor.current);
      const b = ids.indexOf(id);
      if (a < 0 || b < 0) commit([id]);
      else commit(ids.slice(Math.min(a, b), Math.max(a, b) + 1));
    } else {
      commit([id]);
      anchor.current = id;
    }
    setFocusId(id);
  }

  /** Pointer position in the area's scrolled content, kept inside it. */
  function toContent(root: HTMLElement, clientX: number, clientY: number) {
    const box = root.getBoundingClientRect();
    const x = clientX - box.left - root.clientLeft + root.scrollLeft;
    const y = clientY - box.top - root.clientTop + root.scrollTop;
    return {
      x: Math.min(Math.max(x, 0), root.scrollWidth),
      y: Math.min(Math.max(y, 0), root.scrollHeight),
    };
  }

  function sweep() {
    const root = rootRef.current;
    const d = drag.current;
    if (!root || !d) return;
    const p = toContent(root, d.clientX, d.clientY);
    const area = {
      x: Math.min(d.x, p.x),
      y: Math.min(d.y, p.y),
      w: Math.abs(p.x - d.x),
      h: Math.abs(p.y - d.y),
    };
    setMarquee(area);
    const box = root.getBoundingClientRect();
    const hits: string[] = [];
    for (const el of items()) {
      if (el.getAttribute("aria-disabled") === "true") continue;
      const r = el.getBoundingClientRect();
      const left = r.left - box.left - root.clientLeft + root.scrollLeft;
      const top = r.top - box.top - root.clientTop + root.scrollTop;
      if (left < area.x + area.w && left + r.width > area.x && top < area.y + area.h && top + r.height > area.y) {
        hits.push(el.dataset.id ?? "");
      }
    }
    if (d.mode === "add") commit([...d.base, ...hits.filter((id) => !d.base.includes(id))]);
    else if (d.mode === "toggle")
      commit([...d.base.filter((id) => !hits.includes(id)), ...hits.filter((id) => !d.base.includes(id))]);
    else commit(hits);
  }

  function tick() {
    const root = rootRef.current;
    const d = drag.current;
    if (!root || !d) return;
    if (d.moved) {
      const box = root.getBoundingClientRect();
      const push = (pos: number, start: number, end: number) =>
        pos < start + EDGE ? -Math.ceil(((start + EDGE - pos) / EDGE) * SPEED) : pos > end - EDGE ? Math.ceil(((pos - end + EDGE) / EDGE) * SPEED) : 0;
      const dx = push(d.clientX, box.left, box.right);
      const dy = push(d.clientY, box.top, box.bottom);
      const before = root.scrollLeft + root.scrollTop;
      if (dx || dy) root.scrollBy(dx, dy);
      if (root.scrollLeft + root.scrollTop !== before) sweep();
    }
    d.frame = requestAnimationFrame(tick);
  }

  function end() {
    const d = drag.current;
    if (!d) return;
    cancelAnimationFrame(d.frame);
    drag.current = null;
    setMarquee(null);
  }

  function handlePointerDown(event: PointerEvent<HTMLDivElement>) {
    onPointerDown?.(event);
    const root = event.currentTarget;
    if (event.defaultPrevented || event.button !== 0 || event.pointerType === "touch") return;
    if (event.target instanceof Element && event.target.closest(".eq-selection-area-item")) return;
    const box = root.getBoundingClientRect();
    /* A press on a scroll bar is not a marquee. */
    if (event.clientX - box.left - root.clientLeft >= root.clientWidth) return;
    if (event.clientY - box.top - root.clientTop >= root.clientHeight) return;
    event.preventDefault();
    root.setPointerCapture(event.pointerId);
    const start = toContent(root, event.clientX, event.clientY);
    const additive = event.shiftKey;
    const toggle = event.metaKey || event.ctrlKey;
    drag.current = {
      pointerId: event.pointerId,
      ...start,
      clientX: event.clientX,
      clientY: event.clientY,
      base: additive || toggle ? latest.current : [],
      mode: toggle ? "toggle" : additive ? "add" : "replace",
      moved: false,
      frame: 0,
    };
    drag.current.frame = requestAnimationFrame(tick);
  }

  function handlePointerMove(event: PointerEvent<HTMLDivElement>) {
    onPointerMove?.(event);
    const d = drag.current;
    if (!d || d.pointerId !== event.pointerId) return;
    d.clientX = event.clientX;
    d.clientY = event.clientY;
    if (!d.moved) {
      const root = event.currentTarget;
      const p = toContent(root, event.clientX, event.clientY);
      if (Math.hypot(p.x - d.x, p.y - d.y) < THRESHOLD) return;
      d.moved = true;
    }
    sweep();
  }

  function handlePointerUp(event: PointerEvent<HTMLDivElement>) {
    onPointerUp?.(event);
    const d = drag.current;
    if (!d || d.pointerId !== event.pointerId) return;
    /* A plain click on empty space clears; a drag has already set the selection. */
    if (!d.moved && d.mode === "replace") commit([]);
    end();
  }

  function handlePointerCancel(event: PointerEvent<HTMLDivElement>) {
    onPointerCancel?.(event);
    end();
  }

  function handleKeyDown(event: KeyboardEvent<HTMLDivElement>) {
    onKeyDown?.(event);
    if (event.defaultPrevented) return;
    const els = items();
    const current = event.target instanceof Element ? event.target.closest<HTMLElement>(".eq-selection-area-item") : null;
    const id = current?.dataset.id;
    const mod = event.metaKey || event.ctrlKey;
    const { key } = event;
    let handled = true;

    if (key.startsWith("Arrow") && current && id !== undefined) {
      const next = nearest(current, els, key);
      if (next) {
        if (event.shiftKey) {
          const add = [id, next.dataset.id ?? ""].filter(
            (x) => !latest.current.includes(x) && els.find((el) => el.dataset.id === x)?.getAttribute("aria-disabled") !== "true",
          );
          if (add.length) commit([...latest.current, ...add]);
        }
        moveFocus(next);
      }
    } else if (key === "Home") moveFocus(els[0]);
    else if (key === "End") moveFocus(els[els.length - 1]);
    else if (key === " " && current && id !== undefined) {
      if (current.getAttribute("aria-disabled") !== "true") {
        commit(latest.current.includes(id) ? latest.current.filter((x) => x !== id) : [...latest.current, id]);
        anchor.current = id;
      }
    } else if (mod && key.toLowerCase() === "a") commit(enabledIds());
    else if (key === "Escape" && latest.current.length > 0) commit([]);
    else handled = false;

    if (handled) {
      event.preventDefault();
      event.stopPropagation();
    }
  }

  const context: Area = {
    selected: new Set(selected),
    focusId,
    pick,
    focus: setFocusId,
  };

  return (
    <div
      role="listbox"
      aria-multiselectable="true"
      aria-label={label}
      {...rest}
      ref={rootRef}
      className={["eq-selection-area", className].filter(Boolean).join(" ")}
      data-dragging={marquee ? "" : undefined}
      onPointerDown={handlePointerDown}
      onPointerMove={handlePointerMove}
      onPointerUp={handlePointerUp}
      onPointerCancel={handlePointerCancel}
      onLostPointerCapture={end}
      onKeyDown={handleKeyDown}
    >
      <AreaContext.Provider value={context}>{children}</AreaContext.Provider>
      {marquee && (
        <div
          className="eq-selection-area-marquee"
          aria-hidden="true"
          style={{ left: marquee.x, top: marquee.y, width: marquee.w, height: marquee.h }}
        />
      )}
    </div>
  );
}

/** The closest item whose centre lies in the arrow's direction, favouring the same row or column. */
function nearest(from: HTMLElement, els: HTMLElement[], key: string) {
  const a = from.getBoundingClientRect();
  const ax = a.left + a.width / 2;
  const ay = a.top + a.height / 2;
  let best: HTMLElement | undefined;
  let bestScore = Infinity;
  for (const el of els) {
    if (el === from) continue;
    const b = el.getBoundingClientRect();
    const dx = b.left + b.width / 2 - ax;
    const dy = b.top + b.height / 2 - ay;
    const main = key === "ArrowRight" ? dx : key === "ArrowLeft" ? -dx : key === "ArrowDown" ? dy : -dy;
    const cross = key === "ArrowRight" || key === "ArrowLeft" ? Math.abs(dy) : Math.abs(dx);
    if (main <= 1) continue;
    const score = main + cross * 2;
    if (score < bestScore) {
      bestScore = score;
      best = el;
    }
  }
  return best;
}

export function SelectionItem({ id, disabled = false, className, children, onClick, onFocus, ...rest }: SelectionItemProps) {
  const area = useContext(AreaContext);
  const isSelected = area?.selected.has(id) ?? false;
  return (
    <div
      role="option"
      {...rest}
      className={["eq-selection-area-item", className].filter(Boolean).join(" ")}
      data-id={id}
      aria-selected={isSelected}
      aria-disabled={disabled || undefined}
      tabIndex={area?.focusId === id ? 0 : -1}
      onClick={(event) => {
        onClick?.(event);
        if (!event.defaultPrevented && !disabled) area?.pick(id, event);
      }}
      onFocus={(event) => {
        onFocus?.(event);
        area?.focus(id);
      }}
    >
      {children}
      <span className="eq-selection-area-tick" aria-hidden="true">
        <svg viewBox="0 0 16 16" fill="none" stroke="currentColor" strokeWidth="1.5" strokeLinecap="round" strokeLinejoin="round">
          <path d="M4.5 8.2l2.3 2.3 4.7-4.9" pathLength={1} />
        </svg>
      </span>
    </div>
  );
}