CedarCopilot/cedar-mail

import { useCallback, useMemo } from 'react';
import type { AnyExtension, Editor } from '@tiptap/core';
import type { Range } from '@tiptap/core';
import { useQuery, useQueryClient } from '@tanstack/react-query';

import {
  FileLinkNode,
  createFileLinkSuggestionExtension,
  useCachedFileLinkSearch,
} from '@/modules/documents/file-link';
import { ConversationNode } from '@/modules/agentCanvas/extensions/ConversationNode';
import { createConversationMention } from '@/modules/documents/mention/ConversationMention';
import type { ConversationMentionItem } from '@/modules/documents/mention/ConversationMentionList';
import { createEventMention } from '@/modules/documents/mention/EventMention';
import { EventNode } from '@/modules/documents/mention/EventNode';
import type { EventMentionItem } from '@/modules/documents/mention/EventMentionList';
import { useTRPC } from '@/providers/query-provider';

function onConversationSelect(
  conversation: ConversationMentionItem,
  editor: Editor,
  range: Range,
) {
  editor.chain().focus().deleteRange(range).run();
  if (!editor.state.schema.nodes.conversationNode) {
    console.warn('[useRichTextExtensions] conversationNode schema not registered');
    return;
  }
  editor.commands.insertContent({
    type: 'conversationNode',
    attrs: { conversationId: conversation.id },
  });
}

function onEventSelect(event: EventMentionItem, editor: Editor, range: Range) {
  editor.chain().focus().deleteRange(range).run();
  editor.commands.insertContent({
    type: 'eventNode',
    attrs: {
      eventId: event.id,
      eventType: event.eventType,
      title: event.title,
    },
  });
}

export function useRichTextExtensions({
  includeConversationMention = true,
}: { includeConversationMention?: boolean } = {}): AnyExtension[] {
  const trpc = useTRPC();
  const queryClient = useQueryClient();

  useQuery(
    trpc.crm.searchConversationsMinimal.queryOptions(
      { query: '', limit: 30 },
      { staleTime: 5 * 60 * 1000, gcTime: Infinity },
    ),
  );

  const searchFileLinks = useCachedFileLinkSearch();

  const searchConversations = useCallback(
    async (query: string): Promise<ConversationMentionItem[]> => {
      type ConversationMinimal = {
        id: string;
        name: string | null;
        companyName: string | null;
        companyLogoUrl: string | null;
      };
      type SearchResult = { conversations: ConversationMinimal[] };
      const toMentionItems = (rows: ConversationMinimal[]): ConversationMentionItem[] =>
        rows.map((c) => ({
          id: c.id,
          name: c.name,
          companyName: c.companyName,
          companyLogoUrl: c.companyLogoUrl,
        }));

      const trimmed = query.trim();
      const recentOptions = trpc.crm.searchConversationsMinimal.queryOptions(
        { query: '', limit: 30 },
        { staleTime: 5 * 60 * 1000, gcTime: Infinity },
      );

      if (!trimmed) {
        const cachedRecent = queryClient.getQueryData<SearchResult>(recentOptions.queryKey);
        if (cachedRecent?.conversations) return toMentionItems(cachedRecent.conversations);
        const result = (await queryClient.fetchQuery(recentOptions)) as SearchResult;
        return toMentionItems(result.conversations ?? []);
      }

      const searchOptions = trpc.crm.searchConversationsMinimal.queryOptions(
        { query: trimmed, limit: 10 },
        { staleTime: 30_000 },
      );
      const cachedSearch = queryClient.getQueryData<SearchResult>(searchOptions.queryKey);
      if (cachedSearch?.conversations) return toMentionItems(cachedSearch.conversations);

      const cachedRecent = queryClient.getQueryData<SearchResult>(recentOptions.queryKey);
      const recentRows = cachedRecent?.conversations ?? [];
      if (recentRows.length > 0) {
        const q = trimmed.toLowerCase();
        const filtered = recentRows.filter(
          (c) =>
            (c.name?.toLowerCase().includes(q) ?? false) ||
            (c.companyName?.toLowerCase().includes(q) ?? false),
        );
        void queryClient.prefetchQuery(searchOptions);
        if (filtered.length > 0) return toMentionItems(filtered.slice(0, 10));
      }

      const result = (await queryClient.fetchQuery(searchOptions)) as SearchResult;
      return toMentionItems(result.conversations ?? []);
    },
    [queryClient, trpc],
  );

  const searchEvents = useCallback(
    async (query: string): Promise<EventMentionItem[]> => {
      const trimmed = query.trim();
      const options = trpc.crm.searchEventsForMention.queryOptions(
        { query: trimmed, limit: 15 },
        { staleTime: 30_000 },
      );
      const cached = queryClient.getQueryData<{ events: EventMentionItem[] }>(options.queryKey);
      if (cached?.events) return cached.events;
      const result = (await queryClient.fetchQuery(options)) as { events: EventMentionItem[] };
      return result.events ?? [];
    },
    [queryClient, trpc],
  );

  return useMemo<AnyExtension[]>(() => {
    const extensions: AnyExtension[] = [
      FileLinkNode,
      createFileLinkSuggestionExtension({ search: searchFileLinks }),
    ];

    if (includeConversationMention) {
      extensions.push(
        ConversationNode,
        createConversationMention({ search: searchConversations, onSelect: onConversationSelect }),
      );
    }

    extensions.push(
      EventNode,
      createEventMention({ search: searchEvents, onSelect: onEventSelect }),
    );

    return extensions;
  }, [includeConversationMention, searchFileLinks, searchConversations, searchEvents]);
}