Annotate

"use client";

import {
  useEffect,
  useId,
  useRef,
  useState,
  type ComponentProps,
  type KeyboardEvent,
  type MouseEvent,
  type PointerEvent,
  type ReactNode,
  type ToggleEvent,
} from "react";
import { Button } from "./button";
import { Textarea } from "./textarea";
import "./annotate.css";

export type AnnotateNote = { id: string; x: number; y: number; text: string };

export type AnnotateProps = Omit<ComponentProps<"div">, "children" | "defaultValue"> & {
  /** Names the reviewed region, e.g. "Checkout page, desktop". */
  label: string;
  /** Controlled notes. x and y are percentages of the region. */
  notes?: AnnotateNote[];
  /** Starting notes when uncontrolled. */
  defaultNotes?: AnnotateNote[];
  onNotesChange?: (notes: AnnotateNote[]) => void;
  /** When false the region behaves normally and pins are read-only. */
  active?: boolean;
  /** The design under review. */
  children: ReactNode;
};

const GAP = 8;
const DRAG_SLOP = 3;

const clamp = (n: number) => Math.round(Math.min(100, Math.max(0, n)) * 10) / 10;

type Drag = {
  id: string;
  pointerId: number;
  x0: number;
  y0: number;
  x: number;
  y: number;
  moved: boolean;
  rect: DOMRect;
};

export function Annotate({
  label,
  notes: notesProp,
  defaultNotes,
  onNotesChange,
  active = true,
  className,
  children,
  ...rest
}: AnnotateProps) {
  const uid = useId();
  const [internal, setInternal] = useState<AnnotateNote[]>(defaultNotes ?? []);
  const notes = notesProp ?? internal;
  const [editing, setEditing] = useState<string | null>(null);
  const [draft, setDraft] = useState("");
  const [dragPos, setDragPos] = useState<{ id: string; x: number; y: number } | null>(null);
  const layerRef = useRef<HTMLDivElement>(null);
  const editorRef = useRef<HTMLDivElement>(null);
  const textRef = useRef<HTMLTextAreaElement>(null);
  const addRef = useRef<HTMLButtonElement>(null);
  const pins = useRef(new Map<string, HTMLButtonElement>());
  const drag = useRef<Drag | null>(null);
  const swallowClick = useRef(false);
  const fresh = useRef<string | null>(null);
  const seq = useRef(0);
  /** Where focus goes after the next render: a note id, or null for the Add note button. */
  const focusNext = useRef<string | null | undefined>(undefined);

  function commit(next: AnnotateNote[]) {
    if (notesProp === undefined) setInternal(next);
    onNotesChange?.(next);
  }

  /** Notes without the previous blank pin, which is dropped when the next one starts. */
  function withoutFresh() {
    const id = fresh.current;
    fresh.current = null;
    return id ? notes.filter((n) => n.id !== id) : notes;
  }

  function open(note: AnnotateNote) {
    if (fresh.current && fresh.current !== note.id) commit(withoutFresh());
    setDraft(note.text);
    setEditing(note.id);
  }

  function create(x: number, y: number) {
    seq.current += 1;
    const note = { id: `${uid}${seq.current}`, x: clamp(x), y: clamp(y), text: "" };
    commit([...withoutFresh(), note]);
    fresh.current = note.id;
    setDraft("");
    setEditing(note.id);
  }

  function addAtCentre() {
    let x = 50;
    let y = 50;
    while (notes.some((n) => Math.abs(n.x - x) < 2 && Math.abs(n.y - y) < 2) && x < 90) {
      x += 4;
      y += 4;
    }
    create(x, y);
  }

  function remove(id: string) {
    const i = notes.findIndex((n) => n.id === id);
    focusNext.current = notes[i + 1]?.id ?? notes[i - 1]?.id ?? null;
    if (fresh.current === id) fresh.current = null;
    commit(notes.filter((n) => n.id !== id));
    if (editing === id) editorRef.current?.hidePopover();
  }

  function move(id: string, x: number, y: number) {
    commit(notes.map((n) => (n.id === id ? { ...n, x: clamp(x), y: clamp(y) } : n)));
  }

  function save() {
    const text = draft.trim();
    if (!editing || !text) return;
    fresh.current = null;
    commit(notes.map((n) => (n.id === editing ? { ...n, text } : n)));
    editorRef.current?.hidePopover();
  }

  useEffect(() => {
    const target = focusNext.current;
    if (target === undefined) return;
    focusNext.current = undefined;
    ((target && pins.current.get(target)) || addRef.current)?.focus();
  });

  useEffect(() => {
    const editor = editorRef.current;
    const pin = editing ? pins.current.get(editing) : undefined;
    if (!editor || !pin || typeof editor.showPopover !== "function") return;
    if (!editor.matches(":popover-open")) editor.showPopover();
    // Placed by its pin, and placed again when the window resizes or the page scrolls under it.
    function place() {
      if (!editor || !pin) return;
      const r = pin.getBoundingClientRect();
      const w = editor.offsetWidth;
      const h = editor.offsetHeight;
      const vw = document.documentElement.clientWidth;
      const vh = document.documentElement.clientHeight;
      const below = r.bottom + GAP + h <= vh - GAP || r.top - GAP - h < GAP;
      const x = r.left + r.width / 2 - w / 2;
      const y = below ? r.bottom + GAP : r.top - GAP - h;
      editor.style.left = `${Math.max(GAP, Math.min(x, vw - w - GAP))}px`;
      editor.style.top = `${Math.max(GAP, Math.min(y, vh - h - GAP))}px`;
      editor.dataset.side = below ? "bottom" : "top";
    }
    place();
    textRef.current?.focus();
    window.addEventListener("resize", place);
    window.addEventListener("scroll", place, { capture: true, passive: true });
    return () => {
      window.removeEventListener("resize", place);
      window.removeEventListener("scroll", place, { capture: true });
    };
  }, [editing]);

  function onEditorToggle(event: ToggleEvent<HTMLDivElement>) {
    if (event.newState !== "closed" || !editing || event.currentTarget.matches(":popover-open")) return;
    const id = editing;
    if (fresh.current === id) {
      fresh.current = null;
      focusNext.current = null;
      commit(notes.filter((n) => n.id !== id));
    } else if (focusNext.current === undefined) {
      const focused = document.activeElement;
      if (!focused || focused === document.body || event.currentTarget.contains(focused)) focusNext.current = id;
    }
    setEditing(null);
  }

  function onLayerClick(event: MouseEvent<HTMLDivElement>) {
    if (!active || event.target !== event.currentTarget) return;
    const r = event.currentTarget.getBoundingClientRect();
    create(((event.clientX - r.left) / r.width) * 100, ((event.clientY - r.top) / r.height) * 100);
  }

  function onPinPointerDown(event: PointerEvent<HTMLButtonElement>, note: AnnotateNote) {
    const layer = layerRef.current;
    if (!layer || event.button !== 0 || !event.isPrimary) return;
    event.currentTarget.setPointerCapture(event.pointerId);
    swallowClick.current = false;
    drag.current = {
      id: note.id,
      pointerId: event.pointerId,
      x0: event.clientX,
      y0: event.clientY,
      x: note.x,
      y: note.y,
      moved: false,
      rect: layer.getBoundingClientRect(),
    };
  }

  function onPinPointerMove(event: PointerEvent<HTMLButtonElement>) {
    const d = drag.current;
    if (!d || d.pointerId !== event.pointerId) return;
    if (!d.moved && Math.hypot(event.clientX - d.x0, event.clientY - d.y0) < DRAG_SLOP) return;
    d.moved = true;
    d.x = clamp(((event.clientX - d.rect.left) / d.rect.width) * 100);
    d.y = clamp(((event.clientY - d.rect.top) / d.rect.height) * 100);
    setDragPos({ id: d.id, x: d.x, y: d.y });
  }

  function onPinPointerUp(event: PointerEvent<HTMLButtonElement>) {
    const d = drag.current;
    if (!d || d.pointerId !== event.pointerId) return;
    drag.current = null;
    if (d.moved) {
      swallowClick.current = true;
      move(d.id, d.x, d.y);
    }
    setDragPos(null);
  }

  function onPinKeyDown(event: KeyboardEvent<HTMLButtonElement>, note: AnnotateNote) {
    const step = event.shiftKey ? 5 : 1;
    const delta: Record<string, [number, number]> = {
      ArrowLeft: [-step, 0],
      ArrowRight: [step, 0],
      ArrowUp: [0, -step],
      ArrowDown: [0, step],
    };
    if (event.key in delta) {
      const [dx, dy] = delta[event.key];
      move(note.id, note.x + dx, note.y + dy);
    } else if (event.key === "Delete" || event.key === "Backspace") {
      remove(note.id);
    } else {
      return;
    }
    event.preventDefault();
    event.stopPropagation();
  }

  const hintId = `${uid}-hint`;
  const editingIndex = notes.findIndex((n) => n.id === editing);

  return (
    <div
      {...rest}
      role="group"
      aria-label={label}
      className={["eq-annotate", className].filter(Boolean).join(" ")}
      data-active={active || undefined}
    >
      <div className="eq-annotate-region">
        {children}
        <div ref={layerRef} className="eq-annotate-layer" onClick={onLayerClick}>
          {notes.map((note, i) => {
            const pos = dragPos?.id === note.id ? dragPos : note;
            const style = { left: `${pos.x}%`, top: `${pos.y}%` };
            if (!active) {
              return (
                <span key={note.id} className="eq-annotate-pin" style={style} title={note.text} aria-hidden="true">
                  {i + 1}
                </span>
              );
            }
            return (
              <button
                key={note.id}
                ref={(el) => {
                  if (el) pins.current.set(note.id, el);
                  else pins.current.delete(note.id);
                }}
                type="button"
                className="eq-annotate-pin"
                style={style}
                aria-label={note.text ? `Note ${i + 1}: ${note.text}` : `Note ${i + 1}, empty`}
                aria-describedby={hintId}
                aria-expanded={editing === note.id}
                data-dragging={dragPos?.id === note.id || undefined}
                onPointerDown={(event) => onPinPointerDown(event, note)}
                onPointerMove={onPinPointerMove}
                onPointerUp={onPinPointerUp}
                onPointerCancel={() => {
                  drag.current = null;
                  setDragPos(null);
                }}
                onClick={() => {
                  if (swallowClick.current) swallowClick.current = false;
                  else open(note);
                }}
                onKeyDown={(event) => onPinKeyDown(event, note)}
              >
                {i + 1}
              </button>
            );
          })}
        </div>
      </div>

      <div className="eq-annotate-footer">
        {notes.length > 0 ? (
          <ol className="eq-annotate-list">
            {notes.map((note, i) => (
              <li key={note.id} className="eq-annotate-item">
                <span className="eq-annotate-number" aria-hidden="true">
                  {i + 1}
                </span>
                <span className="eq-annotate-text" data-empty={!note.text || undefined}>
                  {note.text || "No text yet"}
                </span>
              </li>
            ))}
          </ol>
        ) : (
          active && <p className="eq-annotate-empty">Click anywhere on the design to add a note.</p>
        )}
        {active && (
          <Button ref={addRef} size="s" variant="outline" className="eq-annotate-add" onClick={addAtCentre}>
            Add note
          </Button>
        )}
      </div>

      {active && (
        <>
          <span id={hintId} className="eq-annotate-hint">
            Arrow keys move it, Enter edits, Delete removes.
          </span>
          <div
            ref={editorRef}
            popover="auto"
            role="dialog"
            aria-label={editingIndex >= 0 ? `Note ${editingIndex + 1}` : "Note"}
            className="eq-annotate-editor"
            onToggle={onEditorToggle}
          >
            <Textarea
              ref={textRef}
              id={`${uid}-text`}
              label={editingIndex >= 0 ? `Note ${editingIndex + 1}` : "Note"}
              rows={3}
              value={draft}
              placeholder="What should change here?"
              onChange={(event) => setDraft(event.target.value)}
              onKeyDown={(event) => {
                if (event.key === "Enter" && (event.metaKey || event.ctrlKey)) {
                  event.preventDefault();
                  save();
                }
              }}
            />
            <div className="eq-annotate-actions">
              <Button size="s" variant="ghost" onClick={() => editing && remove(editing)}>
                Delete
              </Button>
              <Button size="s" disabled={!draft.trim()} onClick={save}>
                Save
              </Button>
            </div>
          </div>
        </>
      )}
    </div>
  );
}