"use client"; import { useQuery, useQueryClient } 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, SearchDocumentTitlesResponse } from "@/contracts/types/document.types"; import { documentsApiService } from "@/lib/apis/documents-api.service"; 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; const MIN_SEARCH_LENGTH = 2; const THROTTLE_MS = 200; /** * Throttle hook - fires immediately, then at most once per interval * Better than debounce for typeahead: user sees results updating as they type */ function useThrottled(value: T, delay = THROTTLE_MS) { const [throttled, setThrottled] = useState(value); const lastExecuted = useRef(Date.now()); const timeoutRef = useRef>(); useEffect(() => { const now = Date.now(); const elapsed = now - lastExecuted.current; if (elapsed >= delay) { // Enough time has passed, update immediately lastExecuted.current = now; setThrottled(value); } else { // Schedule update for remaining time if (timeoutRef.current) { clearTimeout(timeoutRef.current); } timeoutRef.current = setTimeout(() => { lastExecuted.current = Date.now(); setThrottled(value); }, delay - elapsed); } return () => { if (timeoutRef.current) { clearTimeout(timeoutRef.current); } }; }, [value, delay]); return throttled; } export const DocumentMentionPicker = forwardRef< DocumentMentionPickerRef, DocumentMentionPickerProps >(function DocumentMentionPicker( { searchSpaceId, onSelectionChange, onDone, initialSelectedDocuments = [], externalSearch = "" }, ref ) { const queryClient = useQueryClient(); // Use external search with throttle (not debounce) for responsive feel const search = externalSearch; const throttledSearch = useThrottled(search, THROTTLE_MS); 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); // Check if search is long enough const isSearchValid = throttledSearch.trim().length >= MIN_SEARCH_LENGTH; const shouldSearch = throttledSearch.trim().length > 0; // Prefetch first page when picker mounts - results appear instantly useEffect(() => { if (!searchSpaceId) return; const prefetchParams = { search_space_id: searchSpaceId, page: 0, page_size: PAGE_SIZE, }; // Prefetch document titles (user docs) queryClient.prefetchQuery({ queryKey: ["document-titles", prefetchParams], queryFn: () => documentsApiService.searchDocumentTitles({ queryParams: prefetchParams }), staleTime: 60 * 1000, }); // Prefetch SurfSense docs queryClient.prefetchQuery({ queryKey: ["surfsense-docs-mention", "", false], queryFn: () => documentsApiService.getSurfsenseDocs({ queryParams: { page: 0, page_size: PAGE_SIZE }, }), staleTime: 3 * 60 * 1000, }); }, [searchSpaceId, queryClient]); // 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); }, [throttledSearch, searchSpaceId]); // Query params for lightweight title search const titleSearchParams = useMemo( () => ({ search_space_id: searchSpaceId, page: 0, page_size: PAGE_SIZE, ...(isSearchValid ? { title: throttledSearch.trim() } : {}), }), [searchSpaceId, throttledSearch, isSearchValid] ); const surfsenseDocsQueryParams = useMemo(() => { const params: { page: number; page_size: number; title?: string } = { page: 0, page_size: PAGE_SIZE, }; if (isSearchValid) { params.title = throttledSearch.trim(); } return params; }, [throttledSearch, isSearchValid]); // Use the new lightweight endpoint for document title search const { data: titleSearchResults, isLoading: isTitleSearchLoading } = useQuery({ queryKey: ["document-titles", titleSearchParams], queryFn: () => documentsApiService.searchDocumentTitles({ queryParams: titleSearchParams }), staleTime: 60 * 1000, // 1 minute - shorter for fresher results enabled: !!searchSpaceId && currentPage === 0 && (!shouldSearch || isSearchValid), }); // Use query for fetching first page of SurfSense docs const { data: surfsenseDocs, isLoading: isSurfsenseDocsLoading } = useQuery({ queryKey: ["surfsense-docs-mention", throttledSearch, isSearchValid], queryFn: () => documentsApiService.getSurfsenseDocs({ queryParams: surfsenseDocsQueryParams }), staleTime: 3 * 60 * 1000, enabled: !shouldSearch || isSearchValid, }); // 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 from lightweight endpoint if (titleSearchResults?.items) { combinedDocs.push(...titleSearchResults.items); setHasMore(titleSearchResults.has_more); } setAccumulatedDocuments(combinedDocs); } }, [titleSearchResults, surfsenseDocs, currentPage]); // Function to load next page using lightweight endpoint const loadNextPage = useCallback(async () => { if (isLoadingMore || !hasMore) return; const nextPage = currentPage + 1; setIsLoadingMore(true); try { const queryParams = { search_space_id: searchSpaceId, page: nextPage, page_size: PAGE_SIZE, ...(isSearchValid ? { title: throttledSearch.trim() } : {}), }; const response: SearchDocumentTitlesResponse = await documentsApiService.searchDocumentTitles({ 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, throttledSearch, searchSpaceId, isSearchValid]); // 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 = (isTitleSearchLoading || isSurfsenseDocsLoading) && currentPage === 0; // Show hint when search is too short const showSearchHint = shouldSearch && !isSearchValid; // 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 */}
{showSearchHint ? (

Type {MIN_SEARCH_LENGTH - throttledSearch.trim().length} more character {MIN_SEARCH_LENGTH - throttledSearch.trim().length > 1 ? "s" : ""} to search

) : 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 && (
)}
)}
); });