"use client"; import { useQuery } from "@tanstack/react-query"; import { FileText } from "lucide-react"; import { forwardRef, useCallback, useEffect, useImperativeHandle, useMemo, useRef, useState, } from "react"; import { getConnectorIcon } from "@/contracts/enums/connectorIcons"; import type { Document, GetDocumentsResponse } from "@/contracts/types/document.types"; import { documentsApiService } from "@/lib/apis/documents-api.service"; import { cacheKeys } from "@/lib/query-client/cache-keys"; import { cn } from "@/lib/utils"; export interface DocumentMentionPickerRef { selectHighlighted: () => void; moveUp: () => void; moveDown: () => void; } interface DocumentMentionPickerProps { searchSpaceId: number; onSelectionChange: (documents: Pick[]) => void; onDone: () => void; initialSelectedDocuments?: Pick[]; externalSearch?: string; } const PAGE_SIZE = 20; function useDebounced(value: T, delay = 300) { const [debounced, setDebounced] = useState(value); useEffect(() => { const t = setTimeout(() => setDebounced(value), delay); return () => clearTimeout(t); }, [value, delay]); return debounced; } export const DocumentMentionPicker = forwardRef< DocumentMentionPickerRef, DocumentMentionPickerProps >(function DocumentMentionPicker( { searchSpaceId, onSelectionChange, onDone, initialSelectedDocuments = [], externalSearch = "" }, ref ) { // Use external search const search = externalSearch; const debouncedSearch = useDebounced(search, 150); const [highlightedIndex, setHighlightedIndex] = useState(0); const itemRefs = useRef>(new Map()); const scrollContainerRef = useRef(null); // State for pagination const [accumulatedDocuments, setAccumulatedDocuments] = useState< Pick[] >([]); const [currentPage, setCurrentPage] = useState(0); const [hasMore, setHasMore] = useState(false); const [isLoadingMore, setIsLoadingMore] = useState(false); // Reset pagination when search or search space changes // biome-ignore lint/correctness/useExhaustiveDependencies: intentionally reset pagination when search/space changes useEffect(() => { setAccumulatedDocuments([]); setCurrentPage(0); setHasMore(false); setHighlightedIndex(0); }, [debouncedSearch, searchSpaceId]); // Query params for initial fetch (page 0) const fetchQueryParams = useMemo( () => ({ search_space_id: searchSpaceId, page: 0, page_size: PAGE_SIZE, }), [searchSpaceId] ); const searchQueryParams = useMemo(() => { return { search_space_id: searchSpaceId, page: 0, page_size: PAGE_SIZE, title: debouncedSearch, }; }, [debouncedSearch, searchSpaceId]); const surfsenseDocsQueryParams = useMemo(() => { const params: { page: number; page_size: number; title?: string } = { page: 0, page_size: PAGE_SIZE, }; if (debouncedSearch.trim()) { params.title = debouncedSearch; } return params; }, [debouncedSearch]); // Use query for fetching first page of documents const { data: documents, isLoading: isDocumentsLoading } = useQuery({ queryKey: cacheKeys.documents.withQueryParams(fetchQueryParams), queryFn: () => documentsApiService.getDocuments({ queryParams: fetchQueryParams }), staleTime: 3 * 60 * 1000, enabled: !!searchSpaceId && !debouncedSearch.trim() && currentPage === 0, }); // Searching - first page const { data: searchedDocuments, isLoading: isSearchedDocumentsLoading } = useQuery({ queryKey: cacheKeys.documents.withQueryParams(searchQueryParams), queryFn: () => documentsApiService.searchDocuments({ queryParams: searchQueryParams }), staleTime: 3 * 60 * 1000, enabled: !!searchSpaceId && !!debouncedSearch.trim() && currentPage === 0, }); // Use query for fetching first page of SurfSense docs const { data: surfsenseDocs, isLoading: isSurfsenseDocsLoading } = useQuery({ queryKey: ["surfsense-docs-mention", debouncedSearch], queryFn: () => documentsApiService.getSurfsenseDocs({ queryParams: surfsenseDocsQueryParams }), staleTime: 3 * 60 * 1000, }); // Update accumulated documents when first page loads - combine both sources useEffect(() => { if (currentPage === 0) { const combinedDocs: Pick[] = []; // Add SurfSense docs first (they appear at top) if (surfsenseDocs?.items) { for (const doc of surfsenseDocs.items) { combinedDocs.push({ id: doc.id, title: doc.title, document_type: "SURFSENSE_DOCS", }); } } // Add regular documents if (debouncedSearch.trim()) { if (searchedDocuments?.items) { combinedDocs.push(...searchedDocuments.items); setHasMore(searchedDocuments.has_more); } } else { if (documents?.items) { combinedDocs.push(...documents.items); setHasMore(documents.has_more); } } setAccumulatedDocuments(combinedDocs); } }, [documents, searchedDocuments, surfsenseDocs, debouncedSearch, currentPage]); // Function to load next page const loadNextPage = useCallback(async () => { if (isLoadingMore || !hasMore) return; const nextPage = currentPage + 1; setIsLoadingMore(true); try { let response: GetDocumentsResponse; if (debouncedSearch.trim()) { const queryParams = { search_space_id: searchSpaceId, page: nextPage, page_size: PAGE_SIZE, title: debouncedSearch, }; response = await documentsApiService.searchDocuments({ queryParams }); } else { const queryParams = { search_space_id: searchSpaceId, page: nextPage, page_size: PAGE_SIZE, }; response = await documentsApiService.getDocuments({ queryParams }); } setAccumulatedDocuments((prev) => [...prev, ...response.items]); setHasMore(response.has_more); setCurrentPage(nextPage); } catch (error) { console.error("Failed to load next page:", error); } finally { setIsLoadingMore(false); } }, [currentPage, hasMore, isLoadingMore, debouncedSearch, searchSpaceId]); // Infinite scroll handler const handleScroll = useCallback( (e: React.UIEvent) => { const target = e.currentTarget; const scrollBottom = target.scrollHeight - target.scrollTop - target.clientHeight; // Load more when within 50px of bottom if (scrollBottom < 50 && hasMore && !isLoadingMore) { loadNextPage(); } }, [hasMore, isLoadingMore, loadNextPage] ); const actualDocuments = accumulatedDocuments; const actualLoading = ((debouncedSearch.trim() ? isSearchedDocumentsLoading : isDocumentsLoading) || isSurfsenseDocsLoading) && currentPage === 0; // Split documents into SurfSense docs and user docs for grouped rendering const surfsenseDocsList = useMemo( () => actualDocuments.filter((doc) => doc.document_type === "SURFSENSE_DOCS"), [actualDocuments] ); const userDocsList = useMemo( () => actualDocuments.filter((doc) => doc.document_type !== "SURFSENSE_DOCS"), [actualDocuments] ); // Track already selected documents using unique key (document_type:id) to avoid ID collisions const selectedKeys = useMemo( () => new Set(initialSelectedDocuments.map((d) => `${d.document_type}:${d.id}`)), [initialSelectedDocuments] ); // Filter out already selected documents for navigation const selectableDocuments = useMemo( () => actualDocuments.filter((doc) => !selectedKeys.has(`${doc.document_type}:${doc.id}`)), [actualDocuments, selectedKeys] ); const handleSelectDocument = useCallback( (doc: Pick) => { onSelectionChange([...initialSelectedDocuments, doc]); onDone(); }, [initialSelectedDocuments, onSelectionChange, onDone] ); // Scroll highlighted item into view useEffect(() => { const item = itemRefs.current.get(highlightedIndex); if (item) { item.scrollIntoView({ block: "nearest", behavior: "smooth" }); } }, [highlightedIndex]); // Reset highlighted index when external search changes const prevSearchRef = useRef(search); if (prevSearchRef.current !== search) { prevSearchRef.current = search; if (highlightedIndex !== 0) { setHighlightedIndex(0); } } // Expose methods to parent via ref useImperativeHandle( ref, () => ({ selectHighlighted: () => { if (selectableDocuments[highlightedIndex]) { handleSelectDocument(selectableDocuments[highlightedIndex]); } }, moveUp: () => { setHighlightedIndex((prev) => (prev > 0 ? prev - 1 : selectableDocuments.length - 1)); }, moveDown: () => { setHighlightedIndex((prev) => (prev < selectableDocuments.length - 1 ? prev + 1 : 0)); }, }), [selectableDocuments, highlightedIndex, handleSelectDocument] ); // Handle keyboard navigation const handleKeyDown = useCallback( (e: React.KeyboardEvent) => { if (selectableDocuments.length === 0) return; switch (e.key) { case "ArrowDown": e.preventDefault(); setHighlightedIndex((prev) => (prev < selectableDocuments.length - 1 ? prev + 1 : 0)); break; case "ArrowUp": e.preventDefault(); setHighlightedIndex((prev) => (prev > 0 ? prev - 1 : selectableDocuments.length - 1)); break; case "Enter": e.preventDefault(); if (selectableDocuments[highlightedIndex]) { handleSelectDocument(selectableDocuments[highlightedIndex]); } break; case "Escape": e.preventDefault(); onDone(); break; } }, [selectableDocuments, highlightedIndex, handleSelectDocument, onDone] ); return (
{/* Document List - Shows max 5 items on mobile, 7-8 items on desktop */}
{actualLoading ? (
) : actualDocuments.length === 0 ? (

No documents found

) : (
{/* SurfSense Documentation Section */} {surfsenseDocsList.length > 0 && ( <>
SurfSense Docs
{surfsenseDocsList.map((doc) => { const docKey = `${doc.document_type}:${doc.id}`; const isAlreadySelected = selectedKeys.has(docKey); const selectableIndex = selectableDocuments.findIndex( (d) => d.document_type === doc.document_type && d.id === doc.id ); const isHighlighted = !isAlreadySelected && selectableIndex === highlightedIndex; return ( ); })} )} {/* User Documents Section */} {userDocsList.length > 0 && ( <>
Your Documents
{userDocsList.map((doc) => { const docKey = `${doc.document_type}:${doc.id}`; const isAlreadySelected = selectedKeys.has(docKey); const selectableIndex = selectableDocuments.findIndex( (d) => d.document_type === doc.document_type && d.id === doc.id ); const isHighlighted = !isAlreadySelected && selectableIndex === highlightedIndex; return ( ); })} )} {/* Loading indicator for additional pages */} {isLoadingMore && (
)}
)}
); });