tusk/src/hooks/use-snapshots.ts

import { useState, useEffect, useCallback, useRef } from "react";
import { useQuery, useMutation, useQueryClient } from "@tanstack/react-query";
import {
  createSnapshot,
  restoreSnapshot,
  listSnapshots,
  readSnapshotMetadata,
  onSnapshotProgress,
} from "@/lib/tauri";
import type {
  CreateSnapshotParams,
  RestoreSnapshotParams,
  SnapshotProgress,
  SnapshotMetadata,
} from "@/types";

export function useListSnapshots() {
  return useQuery({
    queryKey: ["snapshots"],
    queryFn: listSnapshots,
    staleTime: 30_000,
  });
}

export function useReadSnapshotMetadata() {
  return useMutation({
    mutationFn: (filePath: string) => readSnapshotMetadata(filePath),
  });
}

export function useCreateSnapshot() {
  const [progress, setProgress] = useState<SnapshotProgress | null>(null);
  const snapshotIdRef = useRef<string>("");
  const queryClient = useQueryClient();

  const mutation = useMutation({
    mutationFn: ({
      params,
      snapshotId,
      filePath,
    }: {
      params: CreateSnapshotParams;
      snapshotId: string;
      filePath: string;
    }) => {
      snapshotIdRef.current = snapshotId;
      setProgress(null);
      return createSnapshot(params, snapshotId, filePath);
    },
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ["snapshots"] });
    },
  });

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

  const mutationRef = useRef(mutation);
  useEffect(() => {
    mutationRef.current = mutation;
  });

  const reset = useCallback(() => {
    mutationRef.current.reset();
    setProgress(null);
    snapshotIdRef.current = "";
  }, []);

  return {
    create: mutation.mutate,
    result: mutation.data as SnapshotMetadata | undefined,
    error: mutation.error ? String(mutation.error) : null,
    isCreating: mutation.isPending,
    progress,
    reset,
  };
}

export function useRestoreSnapshot() {
  const [progress, setProgress] = useState<SnapshotProgress | null>(null);
  const snapshotIdRef = useRef<string>("");

  const mutation = useMutation({
    mutationFn: ({
      params,
      snapshotId,
    }: {
      params: RestoreSnapshotParams;
      snapshotId: string;
    }) => {
      snapshotIdRef.current = snapshotId;
      setProgress(null);
      return restoreSnapshot(params, snapshotId);
    },
  });

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

  const mutationRef = useRef(mutation);
  useEffect(() => {
    mutationRef.current = mutation;
  });

  const reset = useCallback(() => {
    mutationRef.current.reset();
    setProgress(null);
    snapshotIdRef.current = "";
  }, []);

  return {
    restore: mutation.mutate,
    rowsRestored: mutation.data as number | undefined,
    error: mutation.error ? String(mutation.error) : null,
    isRestoring: mutation.isPending,
    progress,
    reset,
  };
}