Drag Select

"use client";

import * as React from "react";
import { cn } from "@/lib/utils";

interface DragSelectContextValue {
  isSelected: (id: string) => boolean;
  register: (id: string, el: HTMLElement | null) => void;
  toggle: (id: string, additive: boolean) => void;
  focusItem: (id: string) => void;
  isStatic: boolean;
}

const DragSelectContext = React.createContext<DragSelectContextValue | null>(
  null,
);

export interface DragSelectProps
  extends Omit<React.ComponentProps<"div">, "onChange"> {
  /** Controlled set of selected ids. */
  value?: string[];
  defaultValue?: string[];
  onValueChange?: (ids: string[]) => void;
  /** Disables item press/check motion; selection still updates. */
  static?: boolean;
}

interface Rect {
  left: number;
  top: number;
  width: number;
  height: number;
}

/** Marquee only counts once the pointer has traveled this far, px. */
const DRAG_THRESHOLD = 3;

function setsEqual(a: Set<string>, b: Set<string>) {
  if (a.size !== b.size) return false;
  for (const id of a) if (!b.has(id)) return false;
  return true;
}

/**
 * A marquee selection surface over a grid of DragSelectItem children. Dragging
 * paints a rectangle (pointer-captured and single-pointer guarded, so a drag
 * that leaves the surface keeps tracking) and selects every item it
 * intersects; Shift or Cmd/Ctrl adds to the existing selection instead of
 * replacing it. The marquee is drawn with direct DOM writes — no re-render
 * per pointer move; React state only changes when the hit set changes. A
 * plain click on empty surface clears, Escape clears, Cmd/Ctrl+A selects all.
 * Items are focusable checkboxes that toggle with Space/Enter, and selection
 * counts are announced to a polite live region.
 */
export function DragSelect({
  value,
  defaultValue = [],
  onValueChange,
  static: isStatic = false,
  className,
  children,
  onKeyDown,
  ...props
}: DragSelectProps) {
  const surfaceRef = React.useRef<HTMLDivElement>(null);
  const marqueeRef = React.useRef<HTMLDivElement>(null);
  const items = React.useRef<Map<string, HTMLElement>>(new Map());
  const [internal, setInternal] = React.useState<string[]>(defaultValue);
  const selected = value ?? internal;
  const selectedSet = React.useMemo(() => new Set(selected), [selected]);
  const selectedRef = React.useRef(selectedSet);
  selectedRef.current = selectedSet;

  const drag = React.useRef<{
    pointerId: number;
    startX: number;
    startY: number;
    base: Set<string>;
    additive: boolean;
    moved: boolean;
  } | null>(null);

  const commit = React.useCallback(
    (next: Set<string>) => {
      if (setsEqual(next, selectedRef.current)) return;
      const arr = [...next];
      if (value === undefined) setInternal(arr);
      onValueChange?.(arr);
    },
    [value, onValueChange],
  );

  const register = React.useCallback((id: string, el: HTMLElement | null) => {
    if (el) items.current.set(id, el);
    else items.current.delete(id);
  }, []);

  const toggle = React.useCallback(
    (id: string, additive: boolean) => {
      const next = new Set(additive ? selectedRef.current : []);
      if (selectedRef.current.has(id) && additive) next.delete(id);
      else next.add(id);
      commit(next);
    },
    [commit],
  );

  const focusItem = React.useCallback((id: string) => {
    items.current.get(id)?.focus();
  }, []);

  const intersecting = (rect: Rect): Set<string> => {
    const surfaceRect = surfaceRef.current?.getBoundingClientRect();
    if (!surfaceRect) return new Set();
    const hits = new Set<string>();
    items.current.forEach((el, id) => {
      const r = el.getBoundingClientRect();
      const x = r.left - surfaceRect.left;
      const y = r.top - surfaceRect.top;
      const overlaps =
        x < rect.left + rect.width &&
        x + r.width > rect.left &&
        y < rect.top + rect.height &&
        y + r.height > rect.top;
      if (overlaps) hits.add(id);
    });
    return hits;
  };

  /** Geometry goes straight to the DOM — zero re-renders while painting. */
  const paintMarquee = (rect: Rect | null) => {
    const el = marqueeRef.current;
    if (!el) return;
    if (!rect) {
      el.style.opacity = "0";
      return;
    }
    el.style.opacity = "1";
    el.style.transform = `translate3d(${rect.left}px, ${rect.top}px, 0)`;
    el.style.width = `${rect.width}px`;
    el.style.height = `${rect.height}px`;
  };

  const handlePointerDown = (e: React.PointerEvent<HTMLDivElement>) => {
    // Only start a marquee from empty surface, primary button, one pointer.
    if (e.button !== 0 || drag.current) return;
    if ((e.target as HTMLElement).closest("[data-drag-item]")) return;
    const rect = surfaceRef.current?.getBoundingClientRect();
    if (!rect) return;
    e.currentTarget.setPointerCapture(e.pointerId);
    const additive = e.shiftKey || e.metaKey || e.ctrlKey;
    drag.current = {
      pointerId: e.pointerId,
      startX: e.clientX - rect.left,
      startY: e.clientY - rect.top,
      base: additive ? new Set(selectedRef.current) : new Set(),
      additive,
      moved: false,
    };
  };

  const handlePointerMove = (e: React.PointerEvent<HTMLDivElement>) => {
    const state = drag.current;
    if (!state || e.pointerId !== state.pointerId) return;
    const rect = surfaceRef.current?.getBoundingClientRect();
    if (!rect) return;
    const x = e.clientX - rect.left;
    const y = e.clientY - rect.top;
    if (
      !state.moved &&
      Math.abs(x - state.startX) < DRAG_THRESHOLD &&
      Math.abs(y - state.startY) < DRAG_THRESHOLD
    ) {
      return;
    }
    state.moved = true;
    const box: Rect = {
      left: Math.min(x, state.startX),
      top: Math.min(y, state.startY),
      width: Math.abs(x - state.startX),
      height: Math.abs(y - state.startY),
    };
    paintMarquee(box);
    const next = new Set(state.base);
    intersecting(box).forEach((id) => next.add(id));
    commit(next);
  };

  const endDrag = (e: React.PointerEvent<HTMLDivElement>) => {
    const state = drag.current;
    if (!state || e.pointerId !== state.pointerId) return;
    drag.current = null;
    paintMarquee(null);
    // A motionless click on empty surface clears (unless adding).
    if (!state.moved && !state.additive) commit(new Set());
  };

  const context = React.useMemo<DragSelectContextValue>(
    () => ({
      isSelected: (id) => selectedSet.has(id),
      register,
      toggle,
      focusItem,
      isStatic,
    }),
    [selectedSet, register, toggle, focusItem, isStatic],
  );

  return (
    <DragSelectContext.Provider value={context}>
      <div
        ref={surfaceRef}
        data-slot="drag-select"
        onPointerDown={handlePointerDown}
        onPointerMove={handlePointerMove}
        onPointerUp={endDrag}
        onPointerCancel={endDrag}
        onKeyDown={(e) => {
          onKeyDown?.(e);
          if (e.defaultPrevented) return;
          if (e.key === "Escape" && selectedRef.current.size > 0) {
            commit(new Set());
          } else if ((e.metaKey || e.ctrlKey) && e.key === "a") {
            e.preventDefault();
            commit(new Set(items.current.keys()));
          }
        }}
        className={cn("relative touch-none select-none", className)}
        {...props}
      >
        {children}
        <div
          ref={marqueeRef}
          aria-hidden
          className="pointer-events-none absolute top-0 left-0 z-20 rounded-sm border border-primary bg-primary/10 opacity-0"
        />
        <span aria-live="polite" className="sr-only" role="status">
          {selected.length} of {items.current.size} selected
        </span>
      </div>
    </DragSelectContext.Provider>
  );
}

export interface DragSelectItemProps extends React.ComponentProps<"button"> {
  id: string;
}

export function DragSelectItem({
  id,
  className,
  children,
  ...props
}: DragSelectItemProps) {
  const context = React.useContext(DragSelectContext);
  if (!context) {
    throw new Error("<DragSelectItem> must be used within <DragSelect>");
  }
  const selected = context.isSelected(id);

  return (
    <button
      type="button"
      data-drag-item
      role="checkbox"
      aria-checked={selected}
      ref={(el) => context.register(id, el)}
      onClick={(e) => context.toggle(id, e.shiftKey || e.metaKey || e.ctrlKey)}
      className={cn(
        "relative outline-none transition-[box-shadow,background-color,scale] duration-150 ease-out focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2 focus-visible:ring-offset-background",
        !context.isStatic && "active:scale-[0.98]",
        selected && "ring-2 ring-primary ring-offset-2 ring-offset-background",
        className,
      )}
      {...props}
    >
      {children}
      <span
        aria-hidden
        className={cn(
          "absolute top-1 right-1 flex size-4 items-center justify-center rounded-full bg-primary text-primary-foreground",
          !context.isStatic &&
            "transition-[opacity,scale] duration-(--duration-quick) ease-(--ease-out)",
          selected ? "scale-100 opacity-100" : "scale-75 opacity-0",
        )}
      >
        <svg
          viewBox="0 0 10 10"
          className="size-2.5"
          fill="none"
          stroke="currentColor"
          strokeWidth="1.5"
          strokeLinecap="round"
          strokeLinejoin="round"
        >
          <path d="M2 5.2 4.2 7.4 8 3" />
        </svg>
      </span>
    </button>
  );
}