openrag/frontend/src/app/onboarding/page.tsx

"use client";

import { useRouter } from "next/navigation";
import { Suspense, useEffect, useState } from "react";
import { toast } from "sonner";
import {
  type OnboardingVariables,
  useOnboardingMutation,
} from "@/app/api/mutations/useOnboardingMutation";
import { DoclingHealthBanner, useDoclingHealth } from "@/components/docling-health-banner";
import IBMLogo from "@/components/logo/ibm-logo";
import OllamaLogo from "@/components/logo/ollama-logo";
import OpenAILogo from "@/components/logo/openai-logo";
import { ProtectedRoute } from "@/components/protected-route";
import { Button } from "@/components/ui/button";
import {
  Card,
  CardContent,
  CardFooter,
  CardHeader,
} from "@/components/ui/card";
import { DotPattern } from "@/components/ui/dot-pattern";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import {
  Tooltip,
  TooltipContent,
  TooltipTrigger,
} from "@/components/ui/tooltip";
import { cn } from "@/lib/utils";
import { useGetSettingsQuery } from "../api/queries/useGetSettingsQuery";
import { IBMOnboarding } from "./components/ibm-onboarding";
import { OllamaOnboarding } from "./components/ollama-onboarding";
import { OpenAIOnboarding } from "./components/openai-onboarding";

function OnboardingPage() {
  const { data: settingsDb, isLoading: isSettingsLoading } =
    useGetSettingsQuery();
  const { isHealthy: isDoclingHealthy } = useDoclingHealth();

  const redirect = "/";

  const router = useRouter();

  // Redirect if already authenticated or in no-auth mode
  useEffect(() => {
    if (!isSettingsLoading && settingsDb && settingsDb.edited) {
      router.push(redirect);
    }
  }, [isSettingsLoading, settingsDb, router]);

  const [modelProvider, setModelProvider] = useState<string>("openai");

  const [sampleDataset, setSampleDataset] = useState<boolean>(true);

  const handleSetModelProvider = (provider: string) => {
    setModelProvider(provider);
    setSettings({
      model_provider: provider,
      embedding_model: "",
      llm_model: "",
    });
  };

  const [settings, setSettings] = useState<OnboardingVariables>({
    model_provider: modelProvider,
    embedding_model: "",
    llm_model: "",
  });

  // Mutations
  const onboardingMutation = useOnboardingMutation({
    onSuccess: (data) => {
      console.log("Onboarding completed successfully", data);
      router.push(redirect);
    },
    onError: (error) => {
      toast.error("Failed to complete onboarding", {
        description: error.message,
      });
    },
  });

  const handleComplete = () => {
    if (
      !settings.model_provider ||
      !settings.llm_model ||
      !settings.embedding_model
    ) {
      toast.error("Please complete all required fields");
      return;
    }

    // Prepare onboarding data
    const onboardingData: OnboardingVariables = {
      model_provider: settings.model_provider,
      llm_model: settings.llm_model,
      embedding_model: settings.embedding_model,
      sample_data: sampleDataset,
    };

    // Add API key if available
    if (settings.api_key) {
      onboardingData.api_key = settings.api_key;
    }

    // Add endpoint if available
    if (settings.endpoint) {
      onboardingData.endpoint = settings.endpoint;
    }

    // Add project_id if available
    if (settings.project_id) {
      onboardingData.project_id = settings.project_id;
    }

    onboardingMutation.mutate(onboardingData);
  };

  const isComplete = !!settings.llm_model && !!settings.embedding_model && isDoclingHealthy;

  return (
    <div className="min-h-dvh w-full flex gap-5 flex-col items-center justify-center bg-background relative p-4">
      <DotPattern
        width={24}
        height={24}
        cx={1}
        cy={1}
        cr={1}
        className={cn(
          "[mask-image:linear-gradient(to_bottom,white,transparent,transparent)]",
          "text-input/70",
        )}
      />

      <DoclingHealthBanner className="absolute top-0 left-0 right-0 w-full z-20" />

      <div className="flex flex-col items-center gap-5 min-h-[550px] w-full z-10">
        <div className="flex flex-col items-center justify-center gap-4">
          <h1 className="text-2xl font-medium font-chivo">
            Connect a model provider
          </h1>
        </div>
        <Card className="w-full max-w-[600px]">
          <Tabs
            defaultValue={modelProvider}
            onValueChange={handleSetModelProvider}
          >
            <CardHeader>
              <TabsList>
                <TabsTrigger value="openai">
                  <OpenAILogo className="w-4 h-4" />
                  OpenAI
                </TabsTrigger>
                <TabsTrigger value="watsonx">
                  <IBMLogo className="w-4 h-4" />
                  IBM watsonx.ai
                </TabsTrigger>
                <TabsTrigger value="ollama">
                  <OllamaLogo className="w-4 h-4" />
                  Ollama
                </TabsTrigger>
              </TabsList>
            </CardHeader>
            <CardContent>
              <TabsContent value="openai">
                <OpenAIOnboarding
                  setSettings={setSettings}
                  sampleDataset={sampleDataset}
                  setSampleDataset={setSampleDataset}
                />
              </TabsContent>
              <TabsContent value="watsonx">
                <IBMOnboarding
                  setSettings={setSettings}
                  sampleDataset={sampleDataset}
                  setSampleDataset={setSampleDataset}
                />
              </TabsContent>
              <TabsContent value="ollama">
                <OllamaOnboarding
                  setSettings={setSettings}
                  sampleDataset={sampleDataset}
                  setSampleDataset={setSampleDataset}
                />
              </TabsContent>
            </CardContent>
          </Tabs>
          <CardFooter className="flex justify-end">
            <Tooltip>
              <TooltipTrigger asChild>
                <div>
                  <Button
                    size="sm"
                    onClick={handleComplete}
                    disabled={!isComplete}
                    loading={onboardingMutation.isPending}
                  >
                    <span className="select-none">Complete</span>
                  </Button>
                </div>
              </TooltipTrigger>
              {!isComplete && (
                <TooltipContent>
                  {!!settings.llm_model && !!settings.embedding_model && !isDoclingHealthy
                    ? "docling-serve must be running to continue"
                    : "Please fill in all required fields"}
                </TooltipContent>
              )}
            </Tooltip>
          </CardFooter>
        </Card>
      </div>
    </div>
  );
}

export default function ProtectedOnboardingPage() {
  return (
    <ProtectedRoute>
      <Suspense fallback={<div>Loading onboarding...</div>}>
        <OnboardingPage />
      </Suspense>
    </ProtectedRoute>
  );
}