ubiquitous-invention/apps/web/components/editor/extensions/mention.tsx

"use client";

import { Node, mergeAttributes } from "@tiptap/react";
import { PluginKey } from "@tiptap/pm/state";
import Suggestion, { type SuggestionProps } from "@tiptap/suggestion";
import { ReactRenderer } from "@tiptap/react";
import * as React from "react";

import { cn } from "@/lib/utils";

export type MentionEntityType = "user" | "object";

export type MentionItem = {
  id: string;
  label: string;
  type: MentionEntityType;
};

const MOCK_ITEMS: MentionItem[] = [
  { id: "u1", label: "Alice Chen", type: "user" },
  { id: "u2", label: "Jordan Smith", type: "user" },
  { id: "u3", label: "Sam Rivera", type: "user" },
  { id: "o1", label: "Q1 Roadmap", type: "object" },
  { id: "o2", label: "Design System", type: "object" },
  { id: "o3", label: "Customer Onboarding", type: "object" },
];

export const mentionSuggestionPluginKey = new PluginKey("mentionSuggestion");

function filterItems(query: string): MentionItem[] {
  const q = query.trim().toLowerCase();
  if (!q) return MOCK_ITEMS.slice(0, 8);
  return MOCK_ITEMS.filter((item) =>
    item.label.toLowerCase().includes(q),
  ).slice(0, 8);
}

function positionList(el: HTMLElement, props: SuggestionProps<MentionItem>) {
  const rect = props.clientRect?.();
  if (!rect) return;
  el.style.position = "fixed";
  el.style.left = `${rect.left}px`;
  el.style.top = `${rect.bottom + 8}px`;
  el.style.zIndex = "50";
  el.style.minWidth = `${Math.max(rect.width, 220)}px`;
}

type ListProps = {
  items: MentionItem[];
  command: (item: MentionItem) => void;
  selectedIndex: number;
};

function MentionList({ items, command, selectedIndex }: ListProps) {
  const ref = React.useRef<HTMLDivElement>(null);

  React.useEffect(() => {
    const el = ref.current?.children[selectedIndex] as HTMLElement | undefined;
    el?.scrollIntoView({ block: "nearest" });
  }, [selectedIndex]);

  if (items.length === 0) {
    return (
      <div className="rounded-lg border border-border bg-popover px-3 py-2 text-sm text-muted-foreground shadow-lg">
        No matches
      </div>
    );
  }

  return (
    <div
      ref={ref}
      className="max-h-[min(280px,50vh)] overflow-auto rounded-lg border border-border bg-popover p-1 shadow-lg"
      role="listbox"
    >
      {items.map((item, index) => (
        <button
          key={`${item.type}-${item.id}`}
          className={cn(
            "flex w-full items-center rounded-md px-2 py-1.5 text-left text-sm outline-none transition-colors",
            item.type === "user"
              ? "text-foreground"
              : "text-foreground",
            index === selectedIndex && "bg-accent",
          )}
          type="button"
          role="option"
          aria-selected={index === selectedIndex}
          onClick={() => command(item)}
        >
          <span
            className={cn(
              "mr-2 rounded-md px-1.5 py-0.5 text-xs font-medium",
              item.type === "user"
                ? "bg-primary/15 text-primary dark:bg-primary/25"
                : "bg-teal/15 text-teal dark:bg-teal/25",
            )}
          >
            {item.type === "user" ? "User" : "Object"}
          </span>
          <span className="truncate">@{item.label}</span>
        </button>
      ))}
    </div>
  );
}

export const Mention = Node.create({
  name: "mention",
  group: "inline",
  inline: true,
  atom: true,
  selectable: true,

  addAttributes() {
    return {
      id: {
        default: null as string | null,
        parseHTML: (el) => el.getAttribute("data-id"),
        renderHTML: (attrs) => (attrs.id ? { "data-id": attrs.id } : {}),
      },
      label: {
        default: "",
        parseHTML: (el) =>
          el.textContent?.replace(/^@/, "").trim() ?? "",
        renderHTML: () => ({}),
      },
      type: {
        default: "user" as MentionEntityType,
        parseHTML: (el) =>
          (el.getAttribute("data-mention-type") as MentionEntityType | null) ??
          "user",
        renderHTML: (attrs) => ({
          "data-mention-type": attrs.type ?? "user",
        }),
      },
    };
  },

  parseHTML() {
    return [
      {
        tag: 'span[data-type="mention"]',
      },
    ];
  },

  renderHTML({ node, HTMLAttributes }) {
    const label = String(node.attrs.label ?? "");
    const isUser = node.attrs.type === "user";
    return [
      "span",
      mergeAttributes(HTMLAttributes, {
        "data-type": "mention",
        "data-id": node.attrs.id ?? "",
        "data-mention-type": node.attrs.type ?? "user",
        class: cn(
          "mention-chip rounded-md px-1.5 py-0.5 text-sm font-medium",
          isUser
            ? "bg-primary/15 text-primary dark:bg-primary/25"
            : "bg-teal/15 text-teal dark:bg-teal/25",
        ),
      }),
      `@${label}`,
    ];
  },

  renderText({ node }) {
    return `@${node.attrs.label ?? ""}`;
  },

  addProseMirrorPlugins() {
    const extensionName = this.name;

    let renderer: ReactRenderer | null = null;
    let lastProps: SuggestionProps<MentionItem> | null = null;
    let selectedIndex = 0;

    const updateSelected = (next: number) => {
      if (!lastProps || !renderer) return;
      const len = lastProps.items.length;
      if (len === 0) return;
      selectedIndex = ((next % len) + len) % len;
      renderer.updateProps({
        items: lastProps.items,
        command: lastProps.command,
        selectedIndex,
      });
    };

    return [
      Suggestion<MentionItem, MentionItem>({
        editor: this.editor,
        pluginKey: mentionSuggestionPluginKey,
        char: "@",
        allowSpaces: false,
        items: ({ query }) => filterItems(query),
        command: ({ editor, range, props }) => {
          editor
            .chain()
            .focus()
            .insertContentAt(range, [
              {
                type: extensionName,
                attrs: {
                  id: props.id,
                  label: props.label,
                  type: props.type,
                },
              },
              { type: "text", text: " " },
            ])
            .run();
        },
        render: () => ({
          onStart: (props) => {
            lastProps = props;
            selectedIndex = 0;
            renderer = new ReactRenderer(MentionList, {
              editor: props.editor,
              props: {
                items: props.items,
                command: props.command,
                selectedIndex: 0,
              },
            });
            renderer.element.style.pointerEvents = "auto";
            document.body.appendChild(renderer.element);
            positionList(renderer.element, props);
          },
          onUpdate: (props) => {
            lastProps = props;
            selectedIndex = Math.min(
              selectedIndex,
              Math.max(0, props.items.length - 1),
            );
            renderer?.updateProps({
              items: props.items,
              command: props.command,
              selectedIndex,
            });
            if (renderer) positionList(renderer.element, props);
          },
          onExit: () => {
            renderer?.destroy();
            renderer = null;
            lastProps = null;
            selectedIndex = 0;
          },
          onKeyDown: ({ event }) => {
            if (event.key === "ArrowDown") {
              event.preventDefault();
              updateSelected(selectedIndex + 1);
              return true;
            }
            if (event.key === "ArrowUp") {
              event.preventDefault();
              updateSelected(selectedIndex - 1);
              return true;
            }
            if (event.key === "Enter") {
              event.preventDefault();
              const item = lastProps?.items[selectedIndex];
              if (item && lastProps) {
                lastProps.command(item);
              }
              return true;
            }
            return false;
          },
        }),
      }),
    ];
  },
});