PaperClipAI/ui/src/components/DocumentAnnotationLayer.tsx

import { Profiler, useCallback, useEffect, useId, useLayoutEffect, useMemo, useRef, useState } from "react";
import { AlertTriangle, MessageSquarePlus } from "lucide-react";
import type {
  DocumentAnnotationAnchorState,
  DocumentAnnotationThreadStatus,
} from "@paperclipai/shared";
import { Button } from "@/components/ui/button";
import { cn } from "@/lib/utils";
import {
  buildAnchorFromContainerSelection,
  getContainerTextOffset,
  rangesForNormalizedSpan,
} from "@/lib/document-annotation-selection";
import {
  initializeSelectionDebug,
  isSelectionDebugEnabled,
  recordAnnotationCommit,
  recordCaptureSelection,
  recordMarkdownMutations,
  recordSelectionChange,
} from "@/lib/document-annotation-debug";
import type { DocumentAnnotationAnchorSelector } from "@paperclipai/shared";

export interface AnnotationOverlayThread {
  id: string;
  selectedText: string;
  status: DocumentAnnotationThreadStatus;
  anchorState: DocumentAnnotationAnchorState;
  unreadCount?: number;
}

export interface PendingAnchor {
  selector: DocumentAnnotationAnchorSelector;
  selectedText: string;
}

export interface AnnotationAnchorRect {
  top: number;
  left: number;
  width: number;
  height: number;
}

export interface AnnotationLayerProps {
  containerRef: React.RefObject<HTMLElement | null>;
  markdown: string;
  threads: AnnotationOverlayThread[];
  focusedThreadId: string | null;
  onThreadFocus: (threadId: string, rect: AnnotationAnchorRect) => void;
  /** Tracks the most recently captured pending selection. */
  pendingAnchor: PendingAnchor | null;
  onPendingAnchorChange: (anchor: PendingAnchor | null) => void;
  onRequestComment: (anchor: PendingAnchor, rect: AnnotationAnchorRect) => void;
  /** Publishes refreshed geometry for the currently open popover. */
  onAnchorRectChange?: (rect: AnnotationAnchorRect | null) => void;
  /** Disables the "add comment" affordance when set. */
  newCommentDisabled?: boolean;
  newCommentDisabledReason?: string | null;
  /** Hide resolved highlights even when included in the threads list. */
  hideResolved?: boolean;
  /** Test-only: override window object for layout calculations. */
  testWindow?: { innerWidth: number; innerHeight: number };
  /**
   * When this number changes, re-read the current document selection and emit a
   * pending anchor for the keyboard shortcut path.
   */
  captureSelectionRequestId?: number;
  /**
   * Text of a comment currently being composed. We keep this segment brightly
   * highlighted in the document even after the native browser selection is lost
   * (e.g. once focus moves into the composer textarea).
   */
  pendingHighlightText?: string | null;
}

/** Synthetic thread id used to render the in-progress (pending) comment highlight. */
const PENDING_HIGHLIGHT_THREAD_ID = "__paperclip-pending-annotation__";

interface HighlightRect {
  threadId: string;
  status: DocumentAnnotationThreadStatus;
  anchorState: DocumentAnnotationAnchorState;
  top: number;
  left: number;
  width: number;
  height: number;
  /** True for the last rect of this thread (used to anchor a glyph at the run end). */
  isTail: boolean;
  /** True when this run should render with the brighter focused/pending treatment. */
  focused: boolean;
}

interface ToolbarPosition {
  top: number;
  left: number;
}

type NativeHighlightKind = "open" | "focused" | "stale" | "resolved";

type NativeHighlightRanges = Record<NativeHighlightKind, Range[]>;

type CssHighlight = object;

type HighlightConstructor = new (...ranges: Range[]) => CssHighlight;

type HighlightRegistry = {
  set: (name: string, highlight: CssHighlight) => void;
  delete: (name: string) => void;
};

const NATIVE_HIGHLIGHT_NAMES: Record<NativeHighlightKind, string> = {
  open: "paperclip-doc-annotation-open",
  focused: "paperclip-doc-annotation-focused",
  stale: "paperclip-doc-annotation-stale",
  resolved: "paperclip-doc-annotation-resolved",
};

const nativeHighlightInstances = new Map<string, NativeHighlightRanges>();

function getNativeHighlightApi(): { registry: HighlightRegistry; HighlightCtor: HighlightConstructor } | null {
  const css = (globalThis as { CSS?: { highlights?: HighlightRegistry } }).CSS;
  const HighlightCtor = (globalThis as { Highlight?: HighlightConstructor }).Highlight;
  if (!css?.highlights || typeof HighlightCtor !== "function") return null;
  return { registry: css.highlights, HighlightCtor };
}

function emptyNativeHighlightRanges(): NativeHighlightRanges {
  return {
    open: [],
    focused: [],
    stale: [],
    resolved: [],
  };
}

function syncNativeHighlights(api = getNativeHighlightApi()) {
  if (!api) return;
  for (const kind of Object.keys(NATIVE_HIGHLIGHT_NAMES) as NativeHighlightKind[]) {
    const ranges = Array.from(nativeHighlightInstances.values()).flatMap((entry) => entry[kind]);
    const name = NATIVE_HIGHLIGHT_NAMES[kind];
    if (ranges.length === 0) {
      api.registry.delete(name);
    } else {
      api.registry.set(name, new api.HighlightCtor(...ranges));
    }
  }
}

function setNativeHighlightRanges(instanceId: string, ranges: NativeHighlightRanges) {
  if (!getNativeHighlightApi()) return;
  nativeHighlightInstances.set(instanceId, ranges);
  syncNativeHighlights();
}

function clearNativeHighlightRanges(instanceId: string) {
  if (!nativeHighlightInstances.delete(instanceId)) return;
  syncNativeHighlights();
}

function elementFromNode(node: Node | null | undefined): HTMLElement | null {
  if (!node) return null;
  if (node instanceof HTMLElement) return node;
  const parent = node.parentElement;
  return parent instanceof HTMLElement ? parent : null;
}

function selectionTouchesEditableElement(container: HTMLElement, range: Range) {
  for (const node of [range.startContainer, range.endContainer, range.commonAncestorContainer]) {
    const element = elementFromNode(node);
    if (!element || !container.contains(element)) continue;
    const editableElement = element.closest("input, textarea, select, [contenteditable]");
    if (!(editableElement instanceof HTMLElement)) continue;
    if (editableElement.matches("input, textarea, select")) return true;
    const contentEditableValue = editableElement.getAttribute("contenteditable");
    if (
      editableElement.isContentEditable ||
      (contentEditableValue !== null && contentEditableValue.toLowerCase() !== "false")
    ) {
      return true;
    }
  }
  return false;
}

function intersectRects(a: DOMRect, b: DOMRect): DOMRect | null {
  const left = Math.max(a.left, b.left);
  const top = Math.max(a.top, b.top);
  const right = Math.min(a.right, b.right);
  const bottom = Math.min(a.bottom, b.bottom);
  if (right <= left || bottom <= top) return null;
  return {
    x: left,
    y: top,
    left,
    top,
    right,
    bottom,
    width: right - left,
    height: bottom - top,
    toJSON: () => ({}),
  } as DOMRect;
}

function clipsOverflow(element: HTMLElement) {
  if (element.classList.contains("fold-curtain__content")) return true;
  if (typeof window === "undefined" || typeof window.getComputedStyle !== "function") return false;
  const style = window.getComputedStyle(element);
  return [style.overflow, style.overflowX, style.overflowY].some((value) =>
    value === "hidden" || value === "clip" || value === "auto" || value === "scroll",
  );
}

function visibleClipRectForRange(range: Range, container: HTMLElement): DOMRect | null {
  let clipRect = container.getBoundingClientRect();
  let element = elementFromNode(range.commonAncestorContainer);
  while (element) {
    if (clipsOverflow(element)) {
      const nextClipRect = intersectRects(clipRect, element.getBoundingClientRect());
      if (!nextClipRect) return null;
      clipRect = nextClipRect;
    }
    if (element === container) break;
    element = element.parentElement;
  }
  return clipRect;
}

function nativeHighlightKind(input: {
  focused: boolean;
  stale: boolean;
  resolved: boolean;
}): NativeHighlightKind {
  if (input.resolved) return "resolved";
  if (input.stale) return "stale";
  if (input.focused) return "focused";
  return "open";
}

export function DocumentAnnotationLayer({
  containerRef,
  markdown,
  threads,
  focusedThreadId,
  onThreadFocus,
  pendingAnchor,
  onPendingAnchorChange,
  onRequestComment,
  onAnchorRectChange,
  newCommentDisabled = false,
  newCommentDisabledReason = null,
  hideResolved = true,
  captureSelectionRequestId,
  pendingHighlightText = null,
}: AnnotationLayerProps) {
  const [highlightRects, setHighlightRects] = useState<HighlightRect[]>([]);
  const [hoveredThreadId, setHoveredThreadId] = useState<string | null>(null);
  const [toolbarPosition, setToolbarPosition] = useState<ToolbarPosition | null>(null);
  const overlayRef = useRef<HTMLDivElement | null>(null);
  const lastCaptureSelectionRequestIdRef = useRef<number>(0);
  const lastSelectionRectRef = useRef<AnnotationAnchorRect | null>(null);
  const reactId = useId();
  const nativeHighlightInstanceId = useMemo(
    () => `document-annotation-${reactId.replace(/[^a-zA-Z0-9_-]/g, "")}`,
    [reactId],
  );
  const nativeHighlightsSupported = getNativeHighlightApi() !== null;
  const selectionDebugEnabled = isSelectionDebugEnabled();
  if (selectionDebugEnabled) initializeSelectionDebug();

  const visibleThreads = useMemo(() => {
    if (!hideResolved) return threads;
    return threads.filter((thread) => thread.status !== "resolved" || thread.anchorState === "orphaned" || thread.id === focusedThreadId);
  }, [threads, hideResolved, focusedThreadId]);

  const computeHighlightRects = useCallback(() => {
    const container = containerRef.current;
    const overlay = overlayRef.current;
    if (!container || !overlay) {
      clearNativeHighlightRanges(nativeHighlightInstanceId);
      setHighlightRects([]);
      return;
    }
    const overlayRect = overlay.getBoundingClientRect();
    const next: HighlightRect[] = [];
    const nativeRanges = emptyNativeHighlightRanges();
    const pushRunRects = (run: {
      threadId: string;
      status: DocumentAnnotationThreadStatus;
      anchorState: DocumentAnnotationAnchorState;
      focused: boolean;
      selectedText: string;
      nativeKind: NativeHighlightKind;
    }) => {
      const ranges = rangesForNormalizedSpan({
        container,
        selectedText: run.selectedText,
      });
      const startIndex = next.length;
      for (const range of ranges) {
        const visibleClipRect = visibleClipRectForRange(range, container);
        if (!visibleClipRect) continue;
        let rangeIsVisible = false;
        for (const rect of Array.from(range.getClientRects())) {
          if (rect.width === 0 || rect.height === 0) continue;
          const visibleRect = intersectRects(rect, visibleClipRect);
          if (!visibleRect) continue;
          rangeIsVisible = true;
          next.push({
            threadId: run.threadId,
            status: run.status,
            anchorState: run.anchorState,
            focused: run.focused,
            top: visibleRect.top - overlayRect.top,
            left: visibleRect.left - overlayRect.left,
            width: visibleRect.width,
            height: visibleRect.height,
            isTail: false,
          });
        }
        if (rangeIsVisible) nativeRanges[run.nativeKind].push(range);
      }
      if (next.length > startIndex) {
        next[next.length - 1].isTail = true;
      }
    };
    for (const thread of visibleThreads) {
      if (thread.anchorState === "orphaned") continue;
      const isFocused = thread.id === focusedThreadId;
      const isStale = thread.anchorState === "stale";
      const isResolved = thread.status === "resolved";
      pushRunRects({
        threadId: thread.id,
        status: thread.status,
        anchorState: thread.anchorState,
        focused: isFocused,
        selectedText: thread.selectedText,
        nativeKind: nativeHighlightKind({ focused: isFocused, stale: isStale, resolved: isResolved }),
      });
    }
    // Keep the in-progress (pending) comment selection brightly highlighted so the
    // segment stays anchored in the document while the composer is open.
    if (pendingHighlightText && pendingHighlightText.trim().length > 0) {
      pushRunRects({
        threadId: PENDING_HIGHLIGHT_THREAD_ID,
        status: "open",
        anchorState: "active",
        focused: true,
        selectedText: pendingHighlightText,
        nativeKind: "focused",
      });
    }
    setNativeHighlightRanges(nativeHighlightInstanceId, nativeRanges);
    setHighlightRects(next);
    const activeId = pendingHighlightText ? PENDING_HIGHLIGHT_THREAD_ID : focusedThreadId;
    const activeRects = activeId ? next.filter((rect) => rect.threadId === activeId) : [];
    const activeRect = activeRects.find((rect) => rect.isTail) ?? activeRects[0];
    onAnchorRectChange?.(activeRect ? {
      top: activeRect.top,
      left: activeRect.left,
      width: activeRect.width,
      height: activeRect.height,
    } : null);
  }, [containerRef, focusedThreadId, nativeHighlightInstanceId, onAnchorRectChange, pendingHighlightText, visibleThreads]);

  useLayoutEffect(() => {
    computeHighlightRects();
  }, [computeHighlightRects]);

  useEffect(() => () => clearNativeHighlightRanges(nativeHighlightInstanceId), [nativeHighlightInstanceId]);

  useEffect(() => {
    if (typeof window === "undefined") return;
    const container = containerRef.current;
    const overlay = overlayRef.current;
    let cancelled = false;
    let frame: number | null = null;

    const schedule = () => {
      if (cancelled || frame !== null) return;
      frame = window.requestAnimationFrame(() => {
        frame = null;
        if (!cancelled) computeHighlightRects();
      });
    };

    const handleResizeOrScroll = () => schedule();
    window.addEventListener("resize", handleResizeOrScroll);
    window.addEventListener("scroll", handleResizeOrScroll, true);

    const resizeObserver = typeof window.ResizeObserver === "function"
      ? new window.ResizeObserver(schedule)
      : null;
    if (resizeObserver && container) resizeObserver.observe(container);
    if (resizeObserver && overlay) resizeObserver.observe(overlay);

    const mutationObserver = typeof window.MutationObserver === "function" && container
      ? new window.MutationObserver((mutations) => {
        if (selectionDebugEnabled) {
          const markdownMutations = mutations.filter((mutation) =>
            Boolean(elementFromNode(mutation.target)?.closest(".paperclip-markdown")),
          );
          if (markdownMutations.length > 0) recordMarkdownMutations(markdownMutations.length);
        }
        const onlyLayerMutations = mutations.every((mutation) => {
          const target = elementFromNode(mutation.target);
          return !!target?.closest(".paperclip-doc-annotation-layer, .paperclip-doc-annotation-visual-layer");
        });
        if (!onlyLayerMutations) schedule();
      })
      : null;
    if (mutationObserver && container) {
      mutationObserver.observe(container, {
        childList: true,
        characterData: true,
        subtree: true,
        attributes: true,
        attributeFilter: ["class", "style", "data-state", "open", "hidden", "aria-expanded"],
      });
    }

    schedule();

    return () => {
      cancelled = true;
      if (frame !== null) window.cancelAnimationFrame(frame);
      resizeObserver?.disconnect();
      mutationObserver?.disconnect();
      window.removeEventListener("resize", handleResizeOrScroll);
      window.removeEventListener("scroll", handleResizeOrScroll, true);
    };
  }, [computeHighlightRects, containerRef, selectionDebugEnabled]);

  const captureSelection = useCallback((): PendingAnchor | null => {
    const container = containerRef.current;
    const overlay = overlayRef.current;
    if (!container || !overlay) return null;
    const selection = window.getSelection();
    if (!selection || selection.rangeCount === 0 || selection.isCollapsed) return null;
    const range = selection.getRangeAt(0);
    if (!container.contains(range.commonAncestorContainer)) return null;
    if (selectionTouchesEditableElement(container, range)) return null;
    const containerOffset = getContainerTextOffset(container, range);
    if (!containerOffset) return null;
    const anchor = buildAnchorFromContainerSelection({ markdown, containerOffset });
    if (!anchor) return null;
    const overlayRect = overlay.getBoundingClientRect();
    const rect = range.getBoundingClientRect();
    const top = Math.max(0, rect.top - overlayRect.top - 36);
    const left = Math.max(0, rect.left - overlayRect.left + rect.width / 2 - 80);
    setToolbarPosition({ top, left });
    lastSelectionRectRef.current = {
      top: rect.top - overlayRect.top,
      left: rect.left - overlayRect.left,
      width: rect.width,
      height: rect.height,
    };
    return {
      selector: anchor.selector,
      selectedText: containerOffset.selectedText,
    };
  }, [containerRef, markdown]);

  useEffect(() => {
    if (typeof document === "undefined") return;
    const handleSelectionChange = () => {
      const selection = window.getSelection();
      const range = selection && selection.rangeCount > 0 ? selection.getRangeAt(0) : null;
      const selectionIsActive = Boolean(
        selection && !selection.isCollapsed && range && containerRef.current?.contains(range.commonAncestorContainer),
      );
      if (selectionDebugEnabled) recordSelectionChange(selectionIsActive);
      const captureStartedAt = selectionDebugEnabled ? performance.now() : 0;
      const anchor = captureSelection();
      if (selectionDebugEnabled) {
        recordCaptureSelection(performance.now() - captureStartedAt, Boolean(anchor));
      }
      if (!anchor) {
        onPendingAnchorChange(null);
        setToolbarPosition(null);
        return;
      }
      onPendingAnchorChange(anchor);
    };
    document.addEventListener("selectionchange", handleSelectionChange);
    return () => document.removeEventListener("selectionchange", handleSelectionChange);
  }, [captureSelection, containerRef, onPendingAnchorChange, selectionDebugEnabled]);

  useEffect(() => {
    if (captureSelectionRequestId === undefined) return;
    if (captureSelectionRequestId === 0) return;
    if (lastCaptureSelectionRequestIdRef.current === captureSelectionRequestId) return;
    lastCaptureSelectionRequestIdRef.current = captureSelectionRequestId;
    const anchor = captureSelection();
    if (anchor) {
      onPendingAnchorChange(anchor);
      if (lastSelectionRectRef.current) onRequestComment(anchor, lastSelectionRectRef.current);
    }
  }, [captureSelectionRequestId, captureSelection, onPendingAnchorChange, onRequestComment]);

  const handleAddComment = () => {
    if (pendingAnchor && lastSelectionRectRef.current) onRequestComment(pendingAnchor, lastSelectionRectRef.current);
  };

  const content = (
    <>
      {!nativeHighlightsSupported ? (
        <div className="paperclip-doc-annotation-visual-layer pointer-events-none absolute inset-0 z-0" aria-hidden="true">
          <div className="relative h-full w-full">
            {highlightRects.map((rect, index) => {
              const isFocused = rect.focused;
              const isStale = rect.anchorState === "stale";
              const isResolved = rect.status === "resolved";
              return (
                <span
                  key={`visual-${rect.threadId}-${index}`}
                  data-thread-id={rect.threadId}
                  data-anchor-state={rect.anchorState}
                  data-status={rect.status}
                  data-focused={isFocused || undefined}
                  className={cn(
                    "paperclip-doc-annotation-highlight absolute rounded-none transition-colors",
                    // base box treatment (replaces the previous baseline border)
                    isResolved
                      ? "bg-yellow-100 outline outline-1 outline-dashed outline-offset-0 outline-yellow-700/45 dark:bg-yellow-700 dark:outline-yellow-200/45"
                      : isStale
                        ? "bg-yellow-200 outline outline-2 outline-dashed outline-offset-0 outline-yellow-700/65 dark:bg-yellow-600 dark:outline-yellow-200/70"
                        : isFocused
                          ? "bg-yellow-300 outline outline-2 outline-offset-0 outline-yellow-700/85 shadow-(--shadow-extract-6) dark:bg-yellow-500 dark:outline-yellow-200/85"
                          : "bg-yellow-200 dark:bg-yellow-600",
                  )}
                  style={{
                    top: rect.top,
                    left: rect.left,
                    width: rect.width,
                    height: rect.height,
                  }}
                />
              );
            })}
          </div>
        </div>
      ) : null}
      <div
        className="paperclip-doc-annotation-layer pointer-events-none absolute inset-0 z-(--z-2)"
        aria-hidden="true"
      >
        <div ref={overlayRef} className="relative h-full w-full">
          {highlightRects.map((rect, index) => {
            if (rect.threadId === PENDING_HIGHLIGHT_THREAD_ID) return null;
            const isFocused = rect.focused;
            const isHovered = rect.threadId === hoveredThreadId;
            return (
              <button
                key={`${rect.threadId}-${index}`}
                type="button"
                data-thread-id={rect.threadId}
                data-anchor-state={rect.anchorState}
                data-status={rect.status}
                data-focused={isFocused || undefined}
                data-hovered={isHovered || undefined}
                aria-label="Open annotation thread"
                className={cn(
                  "paperclip-doc-annotation-hit-target pointer-events-auto absolute cursor-pointer rounded-none bg-transparent transition-colors",
                  // Tint the run on hover so it's obvious which highlight you're over.
                  isHovered && "bg-amber-400/40 dark:bg-amber-300/30",
                  isFocused && "ring-1 ring-transparent",
                )}
                style={{
                  top: rect.top,
                  left: rect.left,
                  width: rect.width,
                  height: rect.height,
                }}
                onMouseEnter={() => setHoveredThreadId(rect.threadId)}
                onMouseLeave={() =>
                  setHoveredThreadId((current) => (current === rect.threadId ? null : current))
                }
                onMouseDown={(event) => {
                  event.preventDefault();
                  onThreadFocus(rect.threadId, {
                    top: rect.top,
                    left: rect.left,
                    width: rect.width,
                    height: rect.height,
                  });
                }}
              />
            );
          })}
          {highlightRects.map((rect, index) =>
            rect.isTail && rect.anchorState === "stale" ? (
              <span
                key={`tail-${rect.threadId}-${index}`}
                aria-hidden="true"
                data-thread-id={rect.threadId}
                className="paperclip-doc-annotation-tail pointer-events-none absolute inline-flex items-center justify-center rounded-sm bg-amber-500/95 text-amber-50 shadow-sm dark:bg-amber-500/90 dark:text-amber-50"
                style={{
                  top: rect.top + Math.max(0, rect.height / 2 - 8),
                  left: rect.left + rect.width + 2,
                  width: 16,
                  height: 16,
                }}
                title="Anchor moved — needs review"
              >
                <AlertTriangle className="h-3 w-3" />
              </span>
            ) : null,
          )}
          {pendingAnchor && toolbarPosition ? (
            <div
              data-testid="document-annotation-selection-toolbar"
              role="toolbar"
              aria-label="Selection actions"
              className="paperclip-doc-annotation-selection-toolbar pointer-events-auto absolute z-10 flex items-center gap-1 rounded-md border border-border bg-popover px-1 py-1 shadow-md"
              style={{ top: toolbarPosition.top, left: toolbarPosition.left }}
              onMouseDown={(event) => event.preventDefault()}
            >
              <Button
                type="button"
                size="sm"
                variant="ghost"
                className="h-7 gap-1 px-2 text-xs"
                onClick={handleAddComment}
                disabled={newCommentDisabled}
                title={newCommentDisabled
                  ? newCommentDisabledReason ?? undefined
                  : "Add comment on selection (⌘⇧M)"}
              >
                <MessageSquarePlus className="h-3.5 w-3.5" aria-hidden="true" />
                Comment
              </Button>
            </div>
          ) : null}
        </div>
      </div>
    </>
  );

  return selectionDebugEnabled ? (
    <Profiler id="DocumentAnnotationLayer" onRender={recordAnnotationCommit}>
      {content}
    </Profiler>
  ) : content;
}