openrag/frontend/app/onboarding/_components/onboarding-upload.tsx

import { AnimatePresence, motion } from "motion/react";
import { type ChangeEvent, useEffect, useRef, useState } from "react";
import { toast } from "sonner";
import { useCreateFilter } from "@/app/api/mutations/useCreateFilter";
import { useGetNudgesQuery } from "@/app/api/queries/useGetNudgesQuery";
import { useGetTasksQuery } from "@/app/api/queries/useGetTasksQuery";
import { AnimatedProviderSteps } from "@/app/onboarding/_components/animated-provider-steps";
import { Button } from "@/components/ui/button";
import {
  ONBOARDING_UPLOAD_STEPS_KEY,
  ONBOARDING_USER_DOC_FILTER_ID_KEY,
} from "@/lib/constants";
import { uploadFile } from "@/lib/upload-utils";

interface OnboardingUploadProps {
  onComplete: () => void;
}

const OnboardingUpload = ({ onComplete }: OnboardingUploadProps) => {
  const fileInputRef = useRef<HTMLInputElement>(null);
  const [isUploading, setIsUploading] = useState(false);
  const [currentStep, setCurrentStep] = useState<number | null>(null);
  const [uploadedFilename, setUploadedFilename] = useState<string | null>(null);
  const [shouldCreateFilter, setShouldCreateFilter] = useState(false);
  const [isCreatingFilter, setIsCreatingFilter] = useState(false);

  const createFilterMutation = useCreateFilter();

  const STEP_LIST = [
    "Uploading your document",
    "Generating embeddings",
    "Ingesting document",
    "Processing your document",
  ];

  // Query tasks to track completion
  const { data: tasks } = useGetTasksQuery({
    enabled: currentStep !== null, // Only poll when upload has started
    refetchInterval: currentStep !== null ? 1000 : false, // Poll every 1 second during upload
  });

  const { refetch: refetchNudges } = useGetNudgesQuery(null);

  // Monitor tasks and call onComplete when file processing is done
  useEffect(() => {
    if (currentStep === null || !tasks) {
      return;
    }

    // Check if there are any active tasks (pending, running, or processing)
    const activeTasks = tasks.find(
      (task) =>
        task.status === "pending" ||
        task.status === "running" ||
        task.status === "processing",
    );

    // If no active tasks and we have more than 1 task (initial + new upload), complete it
    if (
      (!activeTasks || (activeTasks.processed_files ?? 0) > 0) &&
      tasks.length > 1
    ) {
      // Set to final step to show "Done"
      setCurrentStep(STEP_LIST.length);

      // Create knowledge filter for uploaded document if requested
      // Guard against race condition: only create if not already creating
      if (shouldCreateFilter && uploadedFilename && !isCreatingFilter) {
        // Reset flags immediately (synchronously) to prevent duplicate creation
        setShouldCreateFilter(false);
        const filename = uploadedFilename;
        setUploadedFilename(null);
        setIsCreatingFilter(true);

        // Get display name from filename (remove extension for cleaner name)
        const displayName = filename.includes(".")
          ? filename.substring(0, filename.lastIndexOf("."))
          : filename;

        const queryData = JSON.stringify({
          query: "",
          filters: {
            data_sources: [filename],
            document_types: ["*"],
            owners: ["*"],
            connector_types: ["*"],
          },
          limit: 10,
          scoreThreshold: 0,
          color: "green",
          icon: "file",
        });

        createFilterMutation
          .mutateAsync({
            name: displayName,
            description: `Filter for ${filename}`,
            queryData: queryData,
          })
          .then((result) => {
            if (result.filter?.id && typeof window !== "undefined") {
              localStorage.setItem(
                ONBOARDING_USER_DOC_FILTER_ID_KEY,
                result.filter.id,
              );
              console.log(
                "Created knowledge filter for uploaded document",
                result.filter.id,
              );
            }
          })
          .catch((error) => {
            console.error("Failed to create knowledge filter:", error);
          })
          .finally(() => {
            setIsCreatingFilter(false);
          });
      }

      // Refetch nudges to get new ones
      refetchNudges();

      // Wait a bit before completing
      setTimeout(() => {
        onComplete();
      }, 1000);
    }
  }, [tasks, currentStep, onComplete, refetchNudges, shouldCreateFilter, uploadedFilename]);

  const resetFileInput = () => {
    if (fileInputRef.current) {
      fileInputRef.current.value = "";
    }
  };

  const handleUploadClick = () => {
    fileInputRef.current?.click();
  };

  const performUpload = async (file: File) => {
    setIsUploading(true);
    try {
      setCurrentStep(0);
      const result = await uploadFile(file, true, true); // Pass createFilter=true
      console.log("Document upload task started successfully");

      // Store filename and createFilter flag in state to create filter after ingestion succeeds
      if (result.createFilter && result.filename) {
        setUploadedFilename(result.filename);
        setShouldCreateFilter(true);
      }

      // Move to processing step - task monitoring will handle completion
      setTimeout(() => {
        setCurrentStep(1);
      }, 1500);
    } catch (error) {
      const errorMessage = error instanceof Error ? error.message : "Upload failed";
      console.error("Upload failed", errorMessage);

      // Show error toast notification
      toast.error("Document upload failed", {
        description: errorMessage,
        duration: 5000,
      });

      // Reset on error
      setCurrentStep(null);
    } finally {
      setIsUploading(false);
    }
  };

  const handleFileChange = async (event: ChangeEvent<HTMLInputElement>) => {
    const selectedFile = event.target.files?.[0];
    if (!selectedFile) {
      resetFileInput();
      return;
    }

    try {
      await performUpload(selectedFile);
    } catch (error) {
      console.error(
        "Unable to prepare file for upload",
        (error as Error).message,
      );
    } finally {
      resetFileInput();
    }
  };

  return (
    <AnimatePresence mode="wait">
      {currentStep === null ? (
        <motion.div
          key="user-ingest"
          initial={{ opacity: 1, y: 0 }}
          exit={{ opacity: 0, y: -24 }}
          transition={{ duration: 0.4, ease: "easeInOut" }}
        >
          <Button
            size="sm"
            variant="outline"
            onClick={handleUploadClick}
            disabled={isUploading}
          >
            <div>{isUploading ? "Uploading..." : "Add a document"}</div>
          </Button>
          <input
            ref={fileInputRef}
            type="file"
            onChange={handleFileChange}
            className="hidden"
            accept=".pdf,.doc,.docx,.txt,.md,.rtf,.odt"
          />
        </motion.div>
      ) : (
        <motion.div
          key="ingest-steps"
          initial={{ opacity: 0, y: 24 }}
          animate={{ opacity: 1, y: 0 }}
          transition={{ duration: 0.4, ease: "easeInOut" }}
        >
          <AnimatedProviderSteps
            currentStep={currentStep}
            setCurrentStep={setCurrentStep}
            isCompleted={false}
            steps={STEP_LIST}
            storageKey={ONBOARDING_UPLOAD_STEPS_KEY}
          />
        </motion.div>
      )}
    </AnimatePresence>
  );
};

export default OnboardingUpload;