openrag/frontend/src/app/onboarding/components/ollama-onboarding.tsx

import { useEffect, useState } from "react";
import { useGetSettingsQuery } from "@/app/api/queries/useGetSettingsQuery";
import { LabelInput } from "@/components/label-input";
import { LabelWrapper } from "@/components/label-wrapper";
import OllamaLogo from "@/components/logo/ollama-logo";
import { useDebouncedValue } from "@/lib/debounce";
import type { OnboardingVariables } from "../../api/mutations/useOnboardingMutation";
import { useGetOllamaModelsQuery } from "../../api/queries/useGetModelsQuery";
import { useModelSelection } from "../hooks/useModelSelection";
import { useUpdateSettings } from "../hooks/useUpdateSettings";
import { AdvancedOnboarding } from "./advanced";
import { ModelSelector } from "./model-selector";

export function OllamaOnboarding({
	setSettings,
	sampleDataset,
	setSampleDataset,
}: {
	setSettings: (settings: OnboardingVariables) => void;
	sampleDataset: boolean;
	setSampleDataset: (dataset: boolean) => void;
}) {
	const { data: settings } = useGetSettingsQuery();
	const [endpoint, setEndpoint] = useState(`http://${settings?.localhost_url ?? "localhost"}:11434`);
	const [showConnecting, setShowConnecting] = useState(false);
	const debouncedEndpoint = useDebouncedValue(endpoint, 500);

	// Fetch models from API when endpoint is provided (debounced)
	const {
		data: modelsData,
		isLoading: isLoadingModels,
		error: modelsError,
	} = useGetOllamaModelsQuery(
		debouncedEndpoint ? { endpoint: debouncedEndpoint } : undefined,
	);

	// Use custom hook for model selection logic
	const {
		languageModel,
		embeddingModel,
		setLanguageModel,
		setEmbeddingModel,
		languageModels,
		embeddingModels,
	} = useModelSelection(modelsData);

	// Handle delayed display of connecting state
	useEffect(() => {
		let timeoutId: NodeJS.Timeout;

		if (debouncedEndpoint && isLoadingModels) {
			timeoutId = setTimeout(() => {
				setShowConnecting(true);
			}, 500);
		} else {
			setShowConnecting(false);
		}

		return () => {
			if (timeoutId) {
				clearTimeout(timeoutId);
			}
		};
	}, [debouncedEndpoint, isLoadingModels]);

	const handleSampleDatasetChange = (dataset: boolean) => {
		setSampleDataset(dataset);
	};

	// Update settings when values change
	useUpdateSettings(
		"ollama",
		{
			endpoint,
			languageModel,
			embeddingModel,
		},
		setSettings,
	);

	// Check validation state based on models query
	const hasConnectionError = debouncedEndpoint && modelsError;
	const hasNoModels =
		modelsData &&
		!modelsData.language_models?.length &&
		!modelsData.embedding_models?.length;

	return (
		<>
			<div className="space-y-4">
				<div className="space-y-1">
					<LabelInput
						label="Ollama Base URL"
						helperText="Base URL of your Ollama server"
						id="api-endpoint"
						required
						placeholder="http://localhost:11434"
						value={endpoint}
						onChange={(e) => setEndpoint(e.target.value)}
					/>
					{showConnecting && (
						<p className="text-mmd text-muted-foreground">
							Connecting to Ollama server...
						</p>
					)}
					{hasConnectionError && (
						<p className="text-mmd text-accent-amber-foreground">
							Can’t reach Ollama at {debouncedEndpoint}. Update the base URL or
							start the server.
						</p>
					)}
					{hasNoModels && (
						<p className="text-mmd text-accent-amber-foreground">
							No models found. Install embedding and agent models on your Ollama
							server.
						</p>
					)}
				</div>
				<LabelWrapper
					label="Embedding model"
					helperText="Model used for knowledge ingest and retrieval"
					id="embedding-model"
					required={true}
				>
					<ModelSelector
						options={embeddingModels}
						icon={<OllamaLogo className="w-4 h-4" />}
						noOptionsPlaceholder={
							isLoadingModels
								? "Loading models..."
								: "No embedding models detected. Install an embedding model to continue."
						}
						value={embeddingModel}
						onValueChange={setEmbeddingModel}
					/>
				</LabelWrapper>
				<LabelWrapper
					label="Language model"
					helperText="Model used for chat"
					id="embedding-model"
					required={true}
				>
					<ModelSelector
						options={languageModels}
						icon={<OllamaLogo className="w-4 h-4" />}
						noOptionsPlaceholder={
							isLoadingModels
								? "Loading models..."
								: "No language models detected. Install a language model to continue."
						}
						value={languageModel}
						onValueChange={setLanguageModel}
					/>
				</LabelWrapper>
			</div>
			<AdvancedOnboarding
				sampleDataset={sampleDataset}
				setSampleDataset={handleSampleDatasetChange}
			/>
		</>
	);
}