tusk/src/hooks/use-data-generator.ts

import { useState, useEffect, useCallback, useRef } from "react";
import { useMutation } from "@tanstack/react-query";
import {
  generateTestDataPreview,
  insertGeneratedData,
  onDataGenProgress,
} from "@/lib/tauri";
import type { GenerateDataParams, DataGenProgress, GeneratedDataPreview } from "@/types";

export function useDataGenerator() {
  const [progress, setProgress] = useState<DataGenProgress | null>(null);
  const genIdRef = useRef<string>("");

  const previewMutation = useMutation({
    mutationFn: ({
      params,
      genId,
    }: {
      params: GenerateDataParams;
      genId: string;
    }) => {
      genIdRef.current = genId;
      setProgress(null);
      return generateTestDataPreview(params, genId);
    },
  });

  const insertMutation = useMutation({
    mutationFn: ({
      connectionId,
      preview,
    }: {
      connectionId: string;
      preview: GeneratedDataPreview;
    }) => insertGeneratedData(connectionId, preview),
  });

  useEffect(() => {
    let mounted = true;
    let unlisten: (() => void) | undefined;
    onDataGenProgress((p) => {
      if (mounted && p.gen_id === genIdRef.current) {
        setProgress(p);
      }
    }).then((fn) => {
      if (mounted) {
        unlisten = fn;
      } else {
        fn();
      }
    });
    return () => {
      mounted = false;
      unlisten?.();
    };
  }, []);

  const previewRef = useRef(previewMutation);
  const insertRef = useRef(insertMutation);
  useEffect(() => {
    previewRef.current = previewMutation;
    insertRef.current = insertMutation;
  });

  const reset = useCallback(() => {
    previewRef.current.reset();
    insertRef.current.reset();
    setProgress(null);
    genIdRef.current = "";
  }, []);

  return {
    generatePreview: previewMutation.mutate,
    preview: previewMutation.data as GeneratedDataPreview | undefined,
    isGenerating: previewMutation.isPending,
    generateError: previewMutation.error ? String(previewMutation.error) : null,
    insertData: insertMutation.mutate,
    insertedRows: insertMutation.data as number | undefined,
    isInserting: insertMutation.isPending,
    insertError: insertMutation.error ? String(insertMutation.error) : null,
    progress,
    reset,
  };
}