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;
}

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;
}

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

/**
 * A marquee selection surface over a grid of DragSelectItem children. Dragging
 * paints a rectangle (pointer-captured, so a drag that leaves the surface keeps
 * tracking) and selects every item it intersects; holding Shift adds to the
 * existing selection instead of replacing it. Items are also focusable and
 * toggle with Space/Enter, and Shift extends there too — fully keyboard
 * operable. The marquee is a decorative layer.
 */
export function DragSelect({
  value,
  defaultValue = [],
  onValueChange,
  className,
  children,
  ...props
}: DragSelectProps) {
  const surfaceRef = 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 [marquee, setMarquee] = React.useState<Rect | null>(null);
  const drag = React.useRef<{
    startX: number;
    startY: number;
    base: Set<string>;
    additive: boolean;
  } | null>(null);

  const commit = React.useCallback(
    (next: Set<string>) => {
      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 ? selectedSet : []);
      if (selectedSet.has(id) && additive) next.delete(id);
      else next.add(id);
      commit(next);
    },
    [selectedSet, 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;
  };

  const handlePointerDown = (e: React.PointerEvent<HTMLDivElement>) => {
    // Only start a marquee from empty surface, primary button.
    if (e.button !== 0) 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 startX = e.clientX - rect.left;
    const startY = e.clientY - rect.top;
    drag.current = {
      startX,
      startY,
      base: e.shiftKey ? new Set(selectedSet) : new Set(),
      additive: e.shiftKey,
    };
    setMarquee({ left: startX, top: startY, width: 0, height: 0 });
  };

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

  const endDrag = () => {
    drag.current = null;
    setMarquee(null);
  };

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

  return (
    <DragSelectContext.Provider value={context}>
      <div
        ref={surfaceRef}
        data-slot="drag-select"
        onPointerDown={handlePointerDown}
        onPointerMove={handlePointerMove}
        onPointerUp={endDrag}
        onPointerCancel={endDrag}
        className={cn("relative touch-none select-none", className)}
        {...props}
      >
        {children}
        {marquee && (
          <div
            aria-hidden
            className="pointer-events-none absolute z-20 rounded-sm border border-primary bg-primary/10"
            style={{
              left: marquee.left,
              top: marquee.top,
              width: marquee.width,
              height: marquee.height,
            }}
          />
        )}
      </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)}
      className={cn(
        "relative outline-none transition-[box-shadow,background-color] duration-150 ease-out focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2 focus-visible:ring-offset-background",
        selected && "ring-2 ring-primary ring-offset-2 ring-offset-background",
        className,
      )}
      {...props}
    >
      {children}
    </button>
  );
}