"use client";

import { ChatInput } from "@llamaindex/chat-ui";
import { useAtom, useAtomValue } from "jotai";
import { Brain, Check, FolderOpen, Minus, Plus, PlusCircle, Zap } from "lucide-react";
import { useParams, useRouter } from "next/navigation";
import React, { Suspense, useCallback, useMemo, useState } from "react";
import { documentTypeCountsAtom } from "@/atoms/documents/document-query.atoms";
import { updateLLMPreferencesMutationAtom } from "@/atoms/llm-config/llm-config-mutation.atoms";
import {
	globalLLMConfigsAtom,
	llmConfigsAtom,
	llmPreferencesAtom,
} from "@/atoms/llm-config/llm-config-query.atoms";
import { DocumentsDataTable } from "@/components/chat/DocumentsDataTable";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import {
	Dialog,
	DialogContent,
	DialogDescription,
	DialogFooter,
	DialogTitle,
	DialogTrigger,
} from "@/components/ui/dialog";
import { Input } from "@/components/ui/input";
import {
	Select,
	SelectContent,
	SelectItem,
	SelectTrigger,
	SelectValue,
} from "@/components/ui/select";
import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip";
import { getConnectorIcon } from "@/contracts/enums/connectorIcons";
import type { Document } from "@/contracts/types/document.types";
import { useSearchSourceConnectors } from "@/hooks/use-search-source-connectors";

const DocumentSelector = React.memo(
	({
		onSelectionChange,
		selectedDocuments = [],
	}: {
		onSelectionChange?: (documents: Document[]) => void;
		selectedDocuments?: Document[];
	}) => {
		const { search_space_id } = useParams();
		const [isOpen, setIsOpen] = useState(false);

		const handleOpenChange = useCallback((open: boolean) => {
			setIsOpen(open);
		}, []);

		const handleSelectionChange = useCallback(
			(documents: Document[]) => {
				onSelectionChange?.(documents);
			},
			[onSelectionChange]
		);

		const handleDone = useCallback(() => {
			setIsOpen(false);
		}, []);

		const selectedCount = React.useMemo(() => selectedDocuments.length, [selectedDocuments.length]);

		return (
			<Dialog open={isOpen} onOpenChange={handleOpenChange}>
				<DialogTrigger asChild>
					<Button
						variant="outline"
						size="sm"
						className="h-9 gap-2 px-3 border-dashed hover:border-solid hover:bg-accent/50 transition-all"
					>
						<FolderOpen className="h-4 w-4 text-muted-foreground" />
						<span className="text-xs font-medium">
							{selectedCount > 0 ? `Selected` : "Documents"}
						</span>
						{selectedCount > 0 && (
							<Badge variant="secondary" className="h-5 px-1.5 text-xs font-medium">
								{selectedCount}
							</Badge>
						)}
					</Button>
				</DialogTrigger>

				<DialogContent className="max-w-[95vw] md:max-w-5xl h-[90vh] md:h-[85vh] p-0 flex flex-col">
					<div className="flex flex-col h-full">
						<div className="px-4 md:px-6 py-4 border-b flex-shrink-0 bg-muted/30">
							<DialogTitle className="text-lg md:text-xl font-semibold">
								Select Documents
							</DialogTitle>
							<DialogDescription className="mt-1.5 text-sm">
								Choose specific documents to include in your research context
							</DialogDescription>
						</div>

						<div className="flex-1 min-h-0 p-4 md:p-6">
							<DocumentsDataTable
								searchSpaceId={Number(search_space_id)}
								onSelectionChange={handleSelectionChange}
								onDone={handleDone}
								initialSelectedDocuments={selectedDocuments}
							/>
						</div>
					</div>
				</DialogContent>
			</Dialog>
		);
	}
);

DocumentSelector.displayName = "DocumentSelector";

const ConnectorSelector = React.memo(
	({
		onSelectionChange,
		selectedConnectors = [],
	}: {
		onSelectionChange?: (connectorTypes: string[]) => void;
		selectedConnectors?: string[];
	}) => {
		const { search_space_id } = useParams();
		const router = useRouter();
		const [isOpen, setIsOpen] = useState(false);

		// Use the documentTypeCountsAtom for fetching document types
		const [documentTypeCountsQuery] = useAtom(documentTypeCountsAtom);
		const {
			data: documentTypeCountsData,
			isLoading,
			refetch: fetchDocumentTypes,
		} = documentTypeCountsQuery;

		// Transform the response into the expected format
		const documentTypes = useMemo(() => {
			if (!documentTypeCountsData) return [];
			return Object.entries(documentTypeCountsData).map(([type, count]) => ({
				type,
				count,
			}));
		}, [documentTypeCountsData]);

		const isLoaded = !!documentTypeCountsData;

		// Fetch live search connectors immediately (non-indexable)
		const {
			connectors: searchConnectors,
			isLoading: connectorsLoading,
			isLoaded: connectorsLoaded,
			fetchConnectors,
		} = useSearchSourceConnectors(false, Number(search_space_id));

		// Filter for non-indexable connectors (live search)
		const liveSearchConnectors = React.useMemo(
			() => searchConnectors.filter((connector) => !connector.is_indexable),
			[searchConnectors]
		);

		const handleOpenChange = useCallback((open: boolean) => {
			setIsOpen(open);
			// Data is already loaded on mount, no need to fetch again
		}, []);

		const handleConnectorToggle = useCallback(
			(connectorType: string) => {
				const isSelected = selectedConnectors.includes(connectorType);
				const newSelection = isSelected
					? selectedConnectors.filter((type) => type !== connectorType)
					: [...selectedConnectors, connectorType];
				onSelectionChange?.(newSelection);
			},
			[selectedConnectors, onSelectionChange]
		);

		const handleSelectAll = useCallback(() => {
			const allTypes = [
				...documentTypes.map((dt) => dt.type),
				...liveSearchConnectors.map((c) => c.connector_type),
			];
			onSelectionChange?.(allTypes);
		}, [documentTypes, liveSearchConnectors, onSelectionChange]);

		const handleClearAll = useCallback(() => {
			onSelectionChange?.([]);
		}, [onSelectionChange]);

		// Get display name for connector type
		const getDisplayName = (type: string) => {
			return type
				.split("_")
				.map((word) => word.charAt(0) + word.slice(1).toLowerCase())
				.join(" ");
		};

		// Get selected document types with their counts
		const selectedDocTypes = documentTypes.filter((dt) => selectedConnectors.includes(dt.type));
		const selectedLiveConnectors = liveSearchConnectors.filter((c) =>
			selectedConnectors.includes(c.connector_type)
		);

		// Total selected count
		const totalSelectedCount = selectedDocTypes.length + selectedLiveConnectors.length;
		const totalAvailableCount = documentTypes.length + liveSearchConnectors.length;

		return (
			<Dialog open={isOpen} onOpenChange={handleOpenChange}>
				<DialogTrigger asChild>
					<Button
						variant="outline"
						size="sm"
						className="relative h-9 gap-2 px-3 border-dashed hover:border-solid hover:bg-accent/50 transition-all"
					>
						<div className="flex items-center gap-1.5">
							{totalSelectedCount > 0 ? (
								<>
									<div className="flex items-center -space-x-2">
										{selectedDocTypes.slice(0, 2).map((docType) => (
											<div
												key={docType.type}
												className="flex h-6 w-6 items-center justify-center rounded-full border-2 border-background bg-muted"
											>
												{getConnectorIcon(docType.type, "h-3 w-3")}
											</div>
										))}
										{selectedLiveConnectors
											.slice(0, 3 - selectedDocTypes.slice(0, 2).length)
											.map((connector) => (
												<div
													key={connector.id}
													className="flex h-6 w-6 items-center justify-center rounded-full border-2 border-background bg-muted"
												>
													{getConnectorIcon(connector.connector_type, "h-3 w-3")}
												</div>
											))}
									</div>
									<span className="text-xs font-medium">
										{totalSelectedCount} {totalSelectedCount === 1 ? "source" : "sources"}
									</span>
								</>
							) : (
								<>
									<Brain className="h-4 w-4 text-muted-foreground" />
									<span className="text-xs font-medium">Sources</span>
								</>
							)}
						</div>
					</Button>
				</DialogTrigger>

				<DialogContent className="sm:max-w-2xl max-h-[85vh] flex flex-col">
					<div className="space-y-4 flex-1 overflow-y-auto pr-2">
						<div>
							<DialogTitle className="text-xl">Select Sources</DialogTitle>
							<DialogDescription className="mt-1.5">
								Choose indexed document types and live search connectors to include in your search
							</DialogDescription>
						</div>

						{isLoading || connectorsLoading ? (
							<div className="flex justify-center py-8">
								<div className="animate-spin h-8 w-8 border-3 border-primary border-t-transparent rounded-full" />
							</div>
						) : totalAvailableCount === 0 ? (
							<div className="flex flex-col items-center justify-center py-12 text-center">
								<div className="rounded-full bg-muted p-4 mb-4">
									<Brain className="h-8 w-8 text-muted-foreground" />
								</div>
								<h4 className="text-sm font-medium mb-1">No sources found</h4>
								<p className="text-xs text-muted-foreground max-w-xs mb-4">
									Add documents or configure search connectors for this search space
								</p>
								<Button
									onClick={() => {
										setIsOpen(false);
										router.push(`/dashboard/${search_space_id}/sources/add`);
									}}
									className="gap-2"
								>
									<PlusCircle className="h-4 w-4" />
									Add Sources
								</Button>
							</div>
						) : (
							<>
								{/* Live Search Connectors Section */}
								{liveSearchConnectors.length > 0 && (
									<div className="space-y-2">
										<div className="flex items-center gap-2 pb-2">
											<Zap className="h-4 w-4 text-primary" />
											<h3 className="text-sm font-semibold">Live Search Connectors</h3>
											<Badge variant="outline" className="text-xs">
												Real-time
											</Badge>
										</div>
										<div className="grid grid-cols-2 gap-3">
											{liveSearchConnectors.map((connector) => {
												const isSelected = selectedConnectors.includes(connector.connector_type);

												return (
													<button
														key={connector.id}
														onClick={() => handleConnectorToggle(connector.connector_type)}
														type="button"
														className={`group relative flex items-center gap-3 p-3 rounded-lg border-2 transition-all ${
															isSelected
																? "border-primary bg-primary/5 shadow-sm"
																: "border-border hover:border-primary/50 hover:bg-accent/50"
														}`}
													>
														<div
															className={`flex h-10 w-10 items-center justify-center rounded-md transition-colors ${
																isSelected ? "bg-primary/10" : "bg-muted group-hover:bg-primary/5"
															}`}
														>
															{getConnectorIcon(
																connector.connector_type,
																`h-5 w-5 ${isSelected ? "text-primary" : "text-muted-foreground group-hover:text-primary"}`
															)}
														</div>
														<div className="flex-1 text-left min-w-0">
															<div className="flex items-center gap-2">
																<p className="text-sm font-medium truncate">{connector.name}</p>
																{isSelected && (
																	<div className="flex h-5 w-5 items-center justify-center rounded-full bg-primary">
																		<Check className="h-3 w-3 text-primary-foreground" />
																	</div>
																)}
															</div>
															<p className="text-xs text-muted-foreground mt-0.5 truncate">
																{getDisplayName(connector.connector_type)}
															</p>
														</div>
													</button>
												);
											})}
										</div>
									</div>
								)}

								{/* Document Types Section */}
								{documentTypes.length > 0 && (
									<div className="space-y-2">
										<div className="flex items-center gap-2 pb-2">
											<FolderOpen className="h-4 w-4 text-primary" />
											<h3 className="text-sm font-semibold">Indexed Document Types</h3>
											<Badge variant="outline" className="text-xs">
												Stored
											</Badge>
										</div>
										<div className="grid grid-cols-2 gap-3">
											{documentTypes.map((docType) => {
												const isSelected = selectedConnectors.includes(docType.type);

												return (
													<button
														key={docType.type}
														onClick={() => handleConnectorToggle(docType.type)}
														type="button"
														className={`group relative flex items-center gap-3 p-3 rounded-lg border-2 transition-all ${
															isSelected
																? "border-primary bg-primary/5 shadow-sm"
																: "border-border hover:border-primary/50 hover:bg-accent/50"
														}`}
													>
														<div
															className={`flex h-10 w-10 items-center justify-center rounded-md transition-colors ${
																isSelected ? "bg-primary/10" : "bg-muted group-hover:bg-primary/5"
															}`}
														>
															{getConnectorIcon(
																docType.type,
																`h-5 w-5 ${isSelected ? "text-primary" : "text-muted-foreground group-hover:text-primary"}`
															)}
														</div>
														<div className="flex-1 text-left min-w-0">
															<div className="flex items-center gap-2">
																<p className="text-sm font-medium truncate">
																	{getDisplayName(docType.type)}
																</p>
																{isSelected && (
																	<div className="flex h-5 w-5 items-center justify-center rounded-full bg-primary">
																		<Check className="h-3 w-3 text-primary-foreground" />
																	</div>
																)}
															</div>
															<p className="text-xs text-muted-foreground mt-0.5">
																{docType.count} {docType.count === 1 ? "document" : "documents"}
															</p>
														</div>
													</button>
												);
											})}
										</div>
									</div>
								)}
							</>
						)}
					</div>

					{totalAvailableCount > 0 && (
						<DialogFooter className="flex flex-row justify-between items-center gap-2 pt-4 border-t">
							<Button
								variant="ghost"
								size="sm"
								onClick={handleClearAll}
								disabled={selectedConnectors.length === 0}
								className="text-xs"
							>
								Clear All
							</Button>
							<Button
								size="sm"
								onClick={handleSelectAll}
								disabled={selectedConnectors.length === totalAvailableCount}
								className="text-xs"
							>
								Select All ({totalAvailableCount})
							</Button>
						</DialogFooter>
					)}
				</DialogContent>
			</Dialog>
		);
	}
);

ConnectorSelector.displayName = "ConnectorSelector";

const TopKSelector = React.memo(
	({ topK = 10, onTopKChange }: { topK?: number; onTopKChange?: (topK: number) => void }) => {
		const MIN_VALUE = 1;
		const MAX_VALUE = 100;

		const handleIncrement = React.useCallback(() => {
			if (topK < MAX_VALUE) {
				onTopKChange?.(topK + 1);
			}
		}, [topK, onTopKChange]);

		const handleDecrement = React.useCallback(() => {
			if (topK > MIN_VALUE) {
				onTopKChange?.(topK - 1);
			}
		}, [topK, onTopKChange]);

		const handleInputChange = React.useCallback(
			(e: React.ChangeEvent<HTMLInputElement>) => {
				const value = e.target.value;
				// Allow empty input for editing
				if (value === "") {
					return;
				}
				const numValue = parseInt(value, 10);
				if (!isNaN(numValue) && numValue >= MIN_VALUE && numValue <= MAX_VALUE) {
					onTopKChange?.(numValue);
				}
			},
			[onTopKChange]
		);

		const handleInputBlur = React.useCallback(
			(e: React.FocusEvent<HTMLInputElement>) => {
				const value = e.target.value;
				if (value === "") {
					// Reset to default if empty
					onTopKChange?.(10);
					return;
				}
				const numValue = parseInt(value, 10);
				if (isNaN(numValue) || numValue < MIN_VALUE) {
					onTopKChange?.(MIN_VALUE);
				} else if (numValue > MAX_VALUE) {
					onTopKChange?.(MAX_VALUE);
				}
			},
			[onTopKChange]
		);

		return (
			<TooltipProvider>
				<Tooltip delayDuration={200}>
					<TooltipTrigger asChild>
						<div className="flex items-center h-8 border rounded-md bg-background hover:bg-accent/50 transition-colors">
							<Button
								type="button"
								variant="ghost"
								size="icon"
								className="h-full w-7 rounded-l-md rounded-r-none hover:bg-accent border-r"
								onClick={handleDecrement}
								disabled={topK <= MIN_VALUE}
							>
								<Minus className="h-3.5 w-3.5" />
							</Button>
							<div className="flex flex-col items-center justify-center px-2 min-w-[60px]">
								<Input
									type="number"
									value={topK}
									onChange={handleInputChange}
									onBlur={handleInputBlur}
									min={MIN_VALUE}
									max={MAX_VALUE}
									className="h-5 w-full px-1 text-center text-sm font-semibold border-0 bg-transparent focus-visible:ring-0 focus-visible:ring-offset-0 [appearance:textfield] [&::-webkit-outer-spin-button]:appearance-none [&::-webkit-inner-spin-button]:appearance-none"
								/>
								<span className="text-[10px] text-muted-foreground leading-none">Results</span>
							</div>
							<Button
								type="button"
								variant="ghost"
								size="icon"
								className="h-full w-7 rounded-r-md rounded-l-none hover:bg-accent border-l"
								onClick={handleIncrement}
								disabled={topK >= MAX_VALUE}
							>
								<Plus className="h-3.5 w-3.5" />
							</Button>
						</div>
					</TooltipTrigger>
					<TooltipContent side="top" className="max-w-xs">
						<div className="space-y-2">
							<p className="text-sm font-semibold">Results per Source</p>
							<p className="text-xs text-muted-foreground leading-relaxed">
								Control how many results to fetch from each data source. Set a higher number to get
								more information, or a lower number for faster, more focused results.
							</p>
							<div className="flex items-center gap-2 text-xs text-muted-foreground pt-1 border-t">
								<span>Recommended: 5-20</span>
								<span>•</span>
								<span>
									Range: {MIN_VALUE}-{MAX_VALUE}
								</span>
							</div>
						</div>
					</TooltipContent>
				</Tooltip>
			</TooltipProvider>
		);
	}
);

TopKSelector.displayName = "TopKSelector";

const LLMSelector = React.memo(() => {
	const { search_space_id } = useParams();
	const searchSpaceId = Number(search_space_id);

	const {
		data: llmConfigs = [],
		isFetching: llmLoading,
		isError: error,
	} = useAtomValue(llmConfigsAtom);
	const {
		data: globalConfigs = [],
		isFetching: globalConfigsLoading,
		isError: globalConfigsError,
	} = useAtomValue(globalLLMConfigsAtom);

	// Replace useLLMPreferences with jotai atoms
	const { data: preferences = {}, isFetching: preferencesLoading } =
		useAtomValue(llmPreferencesAtom);
	const { mutateAsync: updatePreferences } = useAtomValue(updateLLMPreferencesMutationAtom);

	const isLoading = llmLoading || preferencesLoading || globalConfigsLoading;

	// Combine global and custom configs
	const allConfigs = React.useMemo(() => {
		return [...globalConfigs.map((config) => ({ ...config, is_global: true })), ...llmConfigs];
	}, [globalConfigs, llmConfigs]);

	// Memoize the selected config to avoid repeated lookups
	const selectedConfig = React.useMemo(() => {
		if (!preferences.fast_llm_id || !allConfigs.length) return null;
		return allConfigs.find((config) => config.id === preferences.fast_llm_id) || null;
	}, [preferences.fast_llm_id, allConfigs]);

	// Memoize the display value for the trigger
	const displayValue = React.useMemo(() => {
		if (!selectedConfig) return null;
		return (
			<div className="flex items-center gap-1">
				<span className="font-medium text-xs">{selectedConfig.provider}</span>
				<span className="text-muted-foreground">•</span>
				<span className="hidden sm:inline text-muted-foreground text-xs truncate max-w-[60px]">
					{selectedConfig.name}
				</span>
				{"is_global" in selectedConfig && selectedConfig.is_global && (
					<span className="text-xs">🌐</span>
				)}
			</div>
		);
	}, [selectedConfig]);

	const handleValueChange = React.useCallback(
		(value: string) => {
			const llmId = value ? parseInt(value, 10) : undefined;
			updatePreferences({
				search_space_id: searchSpaceId,
				data: { fast_llm_id: llmId },
			});
		},
		[updatePreferences, searchSpaceId]
	);

	// Loading skeleton
	if (isLoading) {
		return (
			<div className="h-8 min-w-[100px] sm:min-w-[120px]">
				<div className="h-8 rounded-md bg-muted animate-pulse flex items-center px-3">
					<div className="w-3 h-3 rounded bg-muted-foreground/20 mr-2" />
					<div className="h-3 w-16 rounded bg-muted-foreground/20" />
				</div>
			</div>
		);
	}

	// Error state
	if (error || globalConfigsError) {
		return (
			<div className="h-8 min-w-[100px] sm:min-w-[120px]">
				<Button
					variant="outline"
					size="sm"
					className="h-8 px-3 text-xs text-destructive border-destructive/50 hover:bg-destructive/10"
					disabled
				>
					<span className="text-xs">Error</span>
				</Button>
			</div>
		);
	}

	return (
		<div className="h-8 min-w-0">
			<Select
				value={preferences.fast_llm_id?.toString() || ""}
				onValueChange={handleValueChange}
				disabled={isLoading}
			>
				<SelectTrigger className="h-8 w-auto min-w-[100px] sm:min-w-[120px] px-3 text-xs border-border bg-background hover:bg-muted/50 transition-colors duration-200 focus:ring-2 focus:ring-primary/20">
					<div className="flex items-center gap-2 min-w-0">
						<Zap className="h-3 w-3 text-primary flex-shrink-0" />
						<SelectValue placeholder="Fast LLM" className="text-xs">
							{displayValue || <span className="text-muted-foreground">Select LLM</span>}
						</SelectValue>
					</div>
				</SelectTrigger>

				<SelectContent align="end" className="w-[300px] max-h-[400px]">
					<div className="px-3 py-2 text-xs font-medium text-muted-foreground border-b bg-muted/30">
						<div className="flex items-center gap-2">
							<Zap className="h-3 w-3" />
							Fast LLM Selection
						</div>
					</div>

					{allConfigs.length === 0 ? (
						<div className="px-4 py-6 text-center">
							<div className="mx-auto w-12 h-12 rounded-full bg-muted flex items-center justify-center mb-3">
								<Brain className="h-5 w-5 text-muted-foreground" />
							</div>
							<h4 className="text-sm font-medium mb-1">No LLM configurations</h4>
							<p className="text-xs text-muted-foreground mb-3">
								Configure AI models to get started
							</p>
							<Button
								variant="outline"
								size="sm"
								className="text-xs"
								onClick={() => window.open("/settings", "_blank")}
							>
								Open Settings
							</Button>
						</div>
					) : (
						<div className="py-1">
							{/* Global Configurations */}
							{globalConfigs.length > 0 && (
								<>
									<div className="px-3 py-1.5 text-xs font-semibold text-muted-foreground">
										Global Configurations
									</div>
									{globalConfigs.map((config) => (
										<SelectItem
											key={config.id}
											value={config.id.toString()}
											className="px-3 py-2 cursor-pointer hover:bg-accent/50 focus:bg-accent"
										>
											<div className="flex items-center justify-between w-full min-w-0">
												<div className="flex items-center gap-3 min-w-0 flex-1">
													<div className="flex h-8 w-8 items-center justify-center rounded-md bg-primary/10 flex-shrink-0">
														<Brain className="h-4 w-4 text-primary" />
													</div>
													<div className="min-w-0 flex-1">
														<div className="flex items-center gap-2 mb-1 flex-wrap">
															<span className="font-medium text-sm truncate">{config.name}</span>
															<Badge
																variant="outline"
																className="text-xs px-1.5 py-0.5 flex-shrink-0"
															>
																{config.provider}
															</Badge>
															<Badge
																variant="secondary"
																className="text-xs px-1.5 py-0.5 flex-shrink-0"
															>
																🌐 Global
															</Badge>
														</div>
														<p className="text-xs text-muted-foreground font-mono truncate">
															{config.model_name}
														</p>
													</div>
												</div>
											</div>
										</SelectItem>
									))}
								</>
							)}

							{/* Custom Configurations */}
							{llmConfigs.length > 0 && (
								<>
									<div className="px-3 py-1.5 text-xs font-semibold text-muted-foreground">
										Your Configurations
									</div>
									{llmConfigs.map((config) => (
										<SelectItem
											key={config.id}
											value={config.id.toString()}
											className="px-3 py-2 cursor-pointer hover:bg-accent/50 focus:bg-accent"
										>
											<div className="flex items-center justify-between w-full min-w-0">
												<div className="flex items-center gap-3 min-w-0 flex-1">
													<div className="flex h-8 w-8 items-center justify-center rounded-md bg-primary/10 flex-shrink-0">
														<Brain className="h-4 w-4 text-primary" />
													</div>
													<div className="min-w-0 flex-1">
														<div className="flex items-center gap-2 mb-1">
															<span className="font-medium text-sm truncate">{config.name}</span>
															<Badge
																variant="outline"
																className="text-xs px-1.5 py-0.5 flex-shrink-0"
															>
																{config.provider}
															</Badge>
														</div>
														<p className="text-xs text-muted-foreground font-mono truncate">
															{config.model_name}
														</p>
													</div>
												</div>
											</div>
										</SelectItem>
									))}
								</>
							)}
						</div>
					)}
				</SelectContent>
			</Select>
		</div>
	);
});

LLMSelector.displayName = "LLMSelector";

const CustomChatInputOptions = React.memo(
	({
		onDocumentSelectionChange,
		selectedDocuments,
		onConnectorSelectionChange,
		selectedConnectors,
		topK,
		onTopKChange,
	}: {
		onDocumentSelectionChange?: (documents: Document[]) => void;
		selectedDocuments?: Document[];
		onConnectorSelectionChange?: (connectorTypes: string[]) => void;
		selectedConnectors?: string[];
		topK?: number;
		onTopKChange?: (topK: number) => void;
	}) => {
		// Memoize the loading fallback to prevent recreation
		const loadingFallback = React.useMemo(
			() => <div className="h-9 w-24 animate-pulse bg-muted/50 rounded-md" />,
			[]
		);

		return (
			<div className="flex flex-wrap gap-2 items-center">
				<div className="flex items-center gap-2">
					<Suspense fallback={loadingFallback}>
						<DocumentSelector
							onSelectionChange={onDocumentSelectionChange}
							selectedDocuments={selectedDocuments}
						/>
					</Suspense>
					<Suspense fallback={loadingFallback}>
						<ConnectorSelector
							onSelectionChange={onConnectorSelectionChange}
							selectedConnectors={selectedConnectors}
						/>
					</Suspense>
				</div>
				<div className="h-4 w-px bg-border hidden sm:block" />
				<TopKSelector topK={topK} onTopKChange={onTopKChange} />
				<div className="h-4 w-px bg-border hidden sm:block" />
				<LLMSelector />
			</div>
		);
	}
);

CustomChatInputOptions.displayName = "CustomChatInputOptions";

export const ChatInputUI = React.memo(
	({
		onDocumentSelectionChange,
		selectedDocuments,
		onConnectorSelectionChange,
		selectedConnectors,
		topK,
		onTopKChange,
	}: {
		onDocumentSelectionChange?: (documents: Document[]) => void;
		selectedDocuments?: Document[];
		onConnectorSelectionChange?: (connectorTypes: string[]) => void;
		selectedConnectors?: string[];
		topK?: number;
		onTopKChange?: (topK: number) => void;
	}) => {
		return (
			<ChatInput className="p-2">
				<ChatInput.Form className="flex gap-2">
					<ChatInput.Field className="flex-1" />
					<ChatInput.Submit />
				</ChatInput.Form>
				<CustomChatInputOptions
					onDocumentSelectionChange={onDocumentSelectionChange}
					selectedDocuments={selectedDocuments}
					onConnectorSelectionChange={onConnectorSelectionChange}
					selectedConnectors={selectedConnectors}
					topK={topK}
					onTopKChange={onTopKChange}
				/>
			</ChatInput>
		);
	}
);

ChatInputUI.displayName = "ChatInputUI";