SurfSense/surfsense_web/components/new-chat/document-mention-picker.tsx

"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<Document, "id" | "title" | "document_type">[]) => void;
	onDone: () => void;
	initialSelectedDocuments?: Pick<Document, "id" | "title" | "document_type">[];
	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<T>(value: T, delay = THROTTLE_MS) {
	const [throttled, setThrottled] = useState(value);
	const lastExecuted = useRef(Date.now());
	const timeoutRef = useRef<ReturnType<typeof setTimeout>>();

	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<Map<number, HTMLButtonElement>>(new Map());
	const scrollContainerRef = useRef<HTMLDivElement>(null);

	// State for pagination
	const [accumulatedDocuments, setAccumulatedDocuments] = useState<
		Pick<Document, "id" | "title" | "document_type">[]
	>([]);
	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<Document, "id" | "title" | "document_type">[] = [];

			// 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<HTMLDivElement>) => {
			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<Document, "id" | "title" | "document_type">) => {
			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 (
		<div
			className="flex flex-col w-[280px] sm:w-[320px] bg-popover rounded-lg"
			onKeyDown={handleKeyDown}
			role="listbox"
			tabIndex={-1}
		>
			{/* Document List - Shows max 5 items on mobile, 7-8 items on desktop */}
			<div
				ref={scrollContainerRef}
				className="max-h-[180px] sm:max-h-[280px] overflow-y-auto"
				onScroll={handleScroll}
			>
				{showSearchHint ? (
					<div className="flex flex-col items-center justify-center py-4 text-center px-4">
						<p className="text-sm text-muted-foreground">
							Type {MIN_SEARCH_LENGTH - throttledSearch.trim().length} more character
							{MIN_SEARCH_LENGTH - throttledSearch.trim().length > 1 ? "s" : ""} to search
						</p>
					</div>
				) : actualLoading ? (
					<div className="flex items-center justify-center py-4">
						<div className="animate-spin h-5 w-5 border-2 border-primary border-t-transparent rounded-full" />
					</div>
				) : actualDocuments.length === 0 ? (
					<div className="flex flex-col items-center justify-center py-4 text-center px-4">
						<FileText className="h-5 w-5 text-muted-foreground/55 mb-1" />
						<p className="text-sm text-muted-foreground">No documents found</p>
					</div>
				) : (
					<div className="py-1">
						{/* SurfSense Documentation Section */}
						{surfsenseDocsList.length > 0 && (
							<>
								<div className="px-3 py-2 text-xs font-bold text-muted-foreground/55">
									SurfSense Docs
								</div>
								{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 (
										<button
											key={docKey}
											ref={(el) => {
												if (el && selectableIndex >= 0) {
													itemRefs.current.set(selectableIndex, el);
												}
											}}
											type="button"
											onClick={() => !isAlreadySelected && handleSelectDocument(doc)}
											onMouseEnter={() => {
												if (!isAlreadySelected && selectableIndex >= 0) {
													setHighlightedIndex(selectableIndex);
												}
											}}
											disabled={isAlreadySelected}
											className={cn(
												"w-full flex items-center gap-2 px-3 py-2 text-left transition-colors",
												isAlreadySelected ? "opacity-50 cursor-not-allowed" : "cursor-pointer",
												isHighlighted && "bg-accent"
											)}
										>
											<span className="shrink-0 text-muted-foreground text-sm">
												{getConnectorIcon(doc.document_type)}
											</span>
											<span className="flex-1 text-sm truncate" title={doc.title}>
												{doc.title}
											</span>
										</button>
									);
								})}
							</>
						)}

						{/* User Documents Section */}
						{userDocsList.length > 0 && (
							<>
								<div className="px-3 py-2 text-xs font-bold text-muted-foreground/55">
									Your Documents
								</div>
								{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 (
										<button
											key={docKey}
											ref={(el) => {
												if (el && selectableIndex >= 0) {
													itemRefs.current.set(selectableIndex, el);
												}
											}}
											type="button"
											onClick={() => !isAlreadySelected && handleSelectDocument(doc)}
											onMouseEnter={() => {
												if (!isAlreadySelected && selectableIndex >= 0) {
													setHighlightedIndex(selectableIndex);
												}
											}}
											disabled={isAlreadySelected}
											className={cn(
												"w-full flex items-center gap-2 px-3 py-2 text-left transition-colors",
												isAlreadySelected ? "opacity-50 cursor-not-allowed" : "cursor-pointer",
												isHighlighted && "bg-accent"
											)}
										>
											<span className="shrink-0 text-muted-foreground text-sm">
												{getConnectorIcon(doc.document_type)}
											</span>
											<span className="flex-1 text-sm truncate" title={doc.title}>
												{doc.title}
											</span>
										</button>
									);
								})}
							</>
						)}

						{/* Loading indicator for additional pages */}
						{isLoadingMore && (
							<div className="flex items-center justify-center py-2">
								<div className="animate-spin h-4 w-4 border-2 border-primary border-t-transparent rounded-full" />
							</div>
						)}
					</div>
				)}
			</div>
		</div>
	);
});