SurfSense/surfsense_web/components/settings/llm-role-manager.tsx

"use client";

import { useAtomValue } from "jotai";
import {
	AlertCircle,
	Bot,
	CircleCheck,
	CircleDashed,
	FileText,
	ImageIcon,
	RefreshCw,
	ScanEye,
} from "lucide-react";
import { useCallback, useEffect, useRef, useState } from "react";
import { toast } from "sonner";
import {
	globalImageGenConfigsAtom,
	imageGenConfigsAtom,
} from "@/atoms/image-gen-config/image-gen-config-query.atoms";
import { updateLLMPreferencesMutationAtom } from "@/atoms/new-llm-config/new-llm-config-mutation.atoms";
import {
	globalNewLLMConfigsAtom,
	llmPreferencesAtom,
	newLLMConfigsAtom,
} from "@/atoms/new-llm-config/new-llm-config-query.atoms";
import {
	globalVisionLLMConfigsAtom,
	visionLLMConfigsAtom,
} from "@/atoms/vision-llm-config/vision-llm-config-query.atoms";
import { Alert, AlertDescription } from "@/components/ui/alert";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { Card, CardContent } from "@/components/ui/card";
import { Label } from "@/components/ui/label";
import {
	Select,
	SelectContent,
	SelectGroup,
	SelectItem,
	SelectLabel,
	SelectTrigger,
	SelectValue,
} from "@/components/ui/select";
import { Skeleton } from "@/components/ui/skeleton";
import { Spinner } from "@/components/ui/spinner";
import { cn } from "@/lib/utils";

const ROLE_DESCRIPTIONS = {
	agent: {
		icon: Bot,
		title: "Agent LLM",
		description: "Primary LLM for chat interactions and agent operations",
		color: "text-muted-foreground",
		bgColor: "bg-muted",
		prefKey: "agent_llm_id" as const,
		configType: "llm" as const,
	},
	document_summary: {
		icon: FileText,
		title: "Document Summary LLM",
		description: "Handles document summarization and research synthesis",
		color: "text-muted-foreground",
		bgColor: "bg-muted",
		prefKey: "document_summary_llm_id" as const,
		configType: "llm" as const,
	},
	image_generation: {
		icon: ImageIcon,
		title: "Image Generation Model",
		description: "Model used for AI image generation (DALL-E, GPT Image, etc.)",
		color: "text-muted-foreground",
		bgColor: "bg-muted",
		prefKey: "image_generation_config_id" as const,
		configType: "image" as const,
	},
	vision: {
		icon: ScanEye,
		title: "Vision LLM",
		description: "Vision-capable model for screenshot analysis and context extraction",
		color: "text-muted-foreground",
		bgColor: "bg-muted",
		prefKey: "vision_llm_config_id" as const,
		configType: "vision" as const,
	},
};

interface LLMRoleManagerProps {
	searchSpaceId: number;
}

export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
	// LLM configs
	const {
		data: newLLMConfigs = [],
		isFetching: configsLoading,
		error: configsError,
		refetch: refreshConfigs,
	} = useAtomValue(newLLMConfigsAtom);
	const {
		data: globalConfigs = [],
		isFetching: globalConfigsLoading,
		error: globalConfigsError,
	} = useAtomValue(globalNewLLMConfigsAtom);

	// Image gen configs
	const {
		data: userImageConfigs = [],
		isFetching: imageConfigsLoading,
		error: imageConfigsError,
	} = useAtomValue(imageGenConfigsAtom);
	const {
		data: globalImageConfigs = [],
		isFetching: globalImageConfigsLoading,
		error: globalImageConfigsError,
	} = useAtomValue(globalImageGenConfigsAtom);

	// Vision LLM configs
	const {
		data: userVisionConfigs = [],
		isFetching: visionConfigsLoading,
		error: visionConfigsError,
	} = useAtomValue(visionLLMConfigsAtom);
	const {
		data: globalVisionConfigs = [],
		isFetching: globalVisionConfigsLoading,
		error: globalVisionConfigsError,
	} = useAtomValue(globalVisionLLMConfigsAtom);

	// Preferences
	const {
		data: preferences = {},
		isFetching: preferencesLoading,
		error: preferencesError,
	} = useAtomValue(llmPreferencesAtom);

	const { mutateAsync: updatePreferences } = useAtomValue(updateLLMPreferencesMutationAtom);

	const [assignments, setAssignments] = useState(() => ({
		agent_llm_id: preferences.agent_llm_id ?? "",
		document_summary_llm_id: preferences.document_summary_llm_id ?? "",
		image_generation_config_id: preferences.image_generation_config_id ?? "",
		vision_llm_config_id: preferences.vision_llm_config_id ?? "",
	}));

	const [savingRole, setSavingRole] = useState<string | null>(null);
	const savingRef = useRef(false);

	useEffect(() => {
		if (!savingRef.current) {
			setAssignments({
				agent_llm_id: preferences.agent_llm_id ?? "",
				document_summary_llm_id: preferences.document_summary_llm_id ?? "",
				image_generation_config_id: preferences.image_generation_config_id ?? "",
				vision_llm_config_id: preferences.vision_llm_config_id ?? "",
			});
		}
	}, [
		preferences?.agent_llm_id,
		preferences?.document_summary_llm_id,
		preferences?.image_generation_config_id,
		preferences?.vision_llm_config_id,
	]);

	const handleRoleAssignment = useCallback(
		async (prefKey: string, configId: string) => {
			const value = configId === "unassigned" ? "" : parseInt(configId);

			setAssignments((prev) => ({ ...prev, [prefKey]: value }));
			setSavingRole(prefKey);
			savingRef.current = true;

			try {
				await updatePreferences({
					search_space_id: searchSpaceId,
					data: { [prefKey]: value || undefined },
				});
				toast.success("Role assignment updated");
			} finally {
				setSavingRole(null);
				savingRef.current = false;
			}
		},
		[updatePreferences, searchSpaceId]
	);

	// Combine global and custom LLM configs
	const allLLMConfigs = [
		...globalConfigs.map((config) => ({ ...config, is_global: true })),
		...newLLMConfigs.filter((config) => config.id && config.id.toString().trim() !== ""),
	];

	// Combine global and custom image gen configs
	const allImageConfigs = [
		...globalImageConfigs.map((config) => ({ ...config, is_global: true })),
		...(userImageConfigs ?? []).filter((config) => config.id && config.id.toString().trim() !== ""),
	];

	// Combine global and custom vision LLM configs
	const allVisionConfigs = [
		...globalVisionConfigs.map((config) => ({ ...config, is_global: true })),
		...(userVisionConfigs ?? []).filter(
			(config) => config.id && config.id.toString().trim() !== ""
		),
	];

	const isLoading =
		configsLoading ||
		preferencesLoading ||
		globalConfigsLoading ||
		imageConfigsLoading ||
		globalImageConfigsLoading ||
		visionConfigsLoading ||
		globalVisionConfigsLoading;
	const hasError =
		configsError ||
		preferencesError ||
		globalConfigsError ||
		imageConfigsError ||
		globalImageConfigsError ||
		visionConfigsError ||
		globalVisionConfigsError;
	const hasAnyConfigs = allLLMConfigs.length > 0 || allImageConfigs.length > 0;

	return (
		<div className="space-y-5 md:space-y-6">
			{/* Header actions */}
			<div className="flex items-center justify-start">
				<Button
					variant="secondary"
					size="sm"
					onClick={() => refreshConfigs()}
					disabled={isLoading}
					className="gap-2"
				>
					<RefreshCw className={cn("h-3.5 w-3.5", isLoading && "animate-spin")} />
					Refresh
				</Button>
			</div>

			{/* Error Alert */}
			{hasError && (
				<div>
					<Alert variant="destructive" className="py-3 md:py-4">
						<AlertCircle className="h-3 w-3 md:h-4 md:w-4 shrink-0" />
						<AlertDescription className="text-xs md:text-sm">
							{(configsError?.message ?? "Failed to load LLM configurations") ||
								(preferencesError?.message ?? "Failed to load preferences") ||
								(globalConfigsError?.message ?? "Failed to load global configurations")}
						</AlertDescription>
					</Alert>
				</div>
			)}

			{/* Loading Skeleton */}
			{isLoading && (
				<div className="grid gap-4 grid-cols-1 lg:grid-cols-2">
					{["skeleton-a", "skeleton-b", "skeleton-c"].map((key) => (
						<Card key={key} className="border-border/60">
							<CardContent className="p-4 md:p-5 space-y-4">
								{/* Header: icon + title + status */}
								<div className="flex items-start justify-between gap-3">
									<div className="flex items-center gap-3 min-w-0">
										<Skeleton className="h-9 w-9 rounded-lg shrink-0" />
										<div className="space-y-1.5 flex-1">
											<Skeleton className="h-4 w-24 md:w-28" />
											<Skeleton className="h-3 w-40 md:w-52" />
										</div>
									</div>
									<Skeleton className="h-4 w-4 rounded-full shrink-0" />
								</div>
								{/* Label */}
								<div className="space-y-1.5">
									<Skeleton className="h-3 w-20" />
									<Skeleton className="h-9 md:h-10 w-full rounded-md" />
								</div>
								{/* Summary block */}
								<div className="rounded-lg border border-border/50 p-3 space-y-2">
									<div className="flex items-center gap-2">
										<Skeleton className="h-3.5 w-3.5 rounded shrink-0" />
										<Skeleton className="h-3.5 w-28" />
									</div>
									<div className="flex items-center gap-1.5">
										<Skeleton className="h-4 w-14 rounded-full" />
										<Skeleton className="h-3 w-24" />
									</div>
								</div>
							</CardContent>
						</Card>
					))}
				</div>
			)}

			{/* No configs warning */}
			{!isLoading && !hasError && !hasAnyConfigs && (
				<Alert variant="destructive" className="py-3 md:py-4">
					<AlertCircle className="h-3 w-3 md:h-4 md:w-4 shrink-0" />
					<AlertDescription className="text-xs md:text-sm">
						No configurations found. Please add at least one LLM provider or image model in the
						respective settings tabs before assigning roles.
					</AlertDescription>
				</Alert>
			)}

			{/* Role Assignment Cards */}
			{!isLoading && !hasError && hasAnyConfigs && (
				<div className="grid gap-4 grid-cols-1 lg:grid-cols-2">
					{Object.entries(ROLE_DESCRIPTIONS).map(([key, role]) => {
						const IconComponent = role.icon;
						const currentAssignment = assignments[role.prefKey as keyof typeof assignments];

						// Pick the right config lists based on role type
						const roleGlobalConfigs =
							role.configType === "image"
								? globalImageConfigs
								: role.configType === "vision"
									? globalVisionConfigs
									: globalConfigs;
						const roleUserConfigs =
							role.configType === "image"
								? (userImageConfigs ?? []).filter((c) => c.id && c.id.toString().trim() !== "")
								: role.configType === "vision"
									? (userVisionConfigs ?? []).filter((c) => c.id && c.id.toString().trim() !== "")
									: newLLMConfigs.filter((c) => c.id && c.id.toString().trim() !== "");
						const roleAllConfigs =
							role.configType === "image"
								? allImageConfigs
								: role.configType === "vision"
									? allVisionConfigs
									: allLLMConfigs;

						const assignedConfig = roleAllConfigs.find((config) => config.id === currentAssignment);
						const isAssigned = !!assignedConfig;

						return (
							<div key={key}>
								<Card className="group relative overflow-hidden transition-all duration-200 border-border/60 hover:shadow-md h-full">
									<CardContent className="p-4 md:p-5 space-y-4">
										{/* Role Header */}
										<div className="flex items-start justify-between gap-3">
											<div className="flex items-center gap-3 min-w-0">
												<div
													className={cn(
														"flex items-center justify-center w-9 h-9 rounded-lg shrink-0",
														role.bgColor
													)}
												>
													<IconComponent className={cn("w-4 h-4", role.color)} />
												</div>
												<div className="min-w-0">
													<h4 className="text-sm font-semibold tracking-tight">{role.title}</h4>
													<p className="text-[11px] text-muted-foreground/70 mt-0.5">
														{role.description}
													</p>
												</div>
											</div>
											{savingRole === role.prefKey ? (
												<Spinner size="sm" className="shrink-0 mt-0.5 text-muted-foreground" />
											) : isAssigned ? (
												<CircleCheck className="w-4 h-4 text-muted-foreground/40 shrink-0 mt-0.5" />
											) : (
												<CircleDashed className="w-4 h-4 text-muted-foreground/40 shrink-0 mt-0.5" />
											)}
										</div>

										{/* Selector */}
										<div className="space-y-1.5">
											<Label className="text-xs font-medium text-muted-foreground">
												Configuration
											</Label>
											<Select
												value={isAssigned ? currentAssignment.toString() : "unassigned"}
												onValueChange={(value) => handleRoleAssignment(role.prefKey, value)}
											>
												<SelectTrigger className="w-full h-9 md:h-10 text-xs md:text-sm">
													<SelectValue placeholder="Select a configuration" />
												</SelectTrigger>
												<SelectContent className="max-w-[calc(100vw-2rem)] select-none">
													<SelectItem
														value="unassigned"
														className="text-xs md:text-sm py-1.5 md:py-2"
													>
														<span className="text-muted-foreground">Unassigned</span>
													</SelectItem>

													{/* Global Configurations */}
													{roleGlobalConfigs.length > 0 && (
														<SelectGroup>
															<SelectLabel className="text-[11px] md:text-xs font-semibold text-muted-foreground px-2 py-1 md:py-1.5">
																Global Configurations
															</SelectLabel>
															{roleGlobalConfigs.map((config) => {
																const isAuto = "is_auto_mode" in config && config.is_auto_mode;
																return (
																	<SelectItem
																		key={config.id}
																		value={config.id.toString()}
																		className="text-xs md:text-sm py-1.5 md:py-2"
																		textValue={config.name}
																	>
																		<div className="flex items-center gap-1 md:gap-1.5 flex-wrap min-w-0">
																			<span className="truncate text-xs md:text-sm">
																				{config.name}
																			</span>
																			{isAuto && (
																				<Badge
																					variant="secondary"
																					className="text-[8px] md:text-[9px] shrink-0 bg-zinc-200 text-zinc-600 dark:bg-zinc-700 dark:text-zinc-300 [[data-slot=select-trigger]_&]:hidden"
																				>
																					Recommended
																				</Badge>
																			)}
																		</div>
																	</SelectItem>
																);
															})}
														</SelectGroup>
													)}

													{/* Custom Configurations */}
													{roleUserConfigs.length > 0 && (
														<SelectGroup>
															<SelectLabel className="text-[11px] md:text-xs font-semibold text-muted-foreground px-2 py-1 md:py-1.5">
																Your Configurations
															</SelectLabel>
															{roleUserConfigs.map((config) => (
																<SelectItem
																	key={config.id}
																	value={config.id.toString()}
																	className="text-xs md:text-sm py-1.5 md:py-2"
																>
																	<div className="flex items-center gap-1 md:gap-1.5 flex-wrap min-w-0">
																		<span className="truncate text-xs md:text-sm">
																			{config.name}
																		</span>
																	</div>
																</SelectItem>
															))}
														</SelectGroup>
													)}
												</SelectContent>
											</Select>
										</div>
									</CardContent>
								</Card>
							</div>
						);
					})}
				</div>
			)}
		</div>
	);
}