tusk/src/components/snapshots/RestoreSnapshotDialog.tsx

import { useState } from "react";
import {
  Dialog,
  DialogContent,
  DialogHeader,
  DialogTitle,
  DialogFooter,
} from "@/components/ui/dialog";
import { Button } from "@/components/ui/button";
import { Badge } from "@/components/ui/badge";
import { useRestoreSnapshot, useReadSnapshotMetadata } from "@/hooks/use-snapshots";
import { toast } from "sonner";
import { open as openFile } from "@tauri-apps/plugin-dialog";
import {
  Loader2,
  CheckCircle2,
  XCircle,
  Upload,
  AlertTriangle,
  FileJson,
} from "lucide-react";
import type { SnapshotMetadata } from "@/types";

interface Props {
  open: boolean;
  onOpenChange: (open: boolean) => void;
  connectionId: string;
}

type Step = "select" | "confirm" | "progress" | "done";

export function RestoreSnapshotDialog({ open, onOpenChange, connectionId }: Props) {
  const [step, setStep] = useState<Step>("select");
  const [filePath, setFilePath] = useState<string | null>(null);
  const [metadata, setMetadata] = useState<SnapshotMetadata | null>(null);
  const [truncate, setTruncate] = useState(false);

  const readMeta = useReadSnapshotMetadata();
  const { restore, rowsRestored, error, isRestoring, progress, reset } = useRestoreSnapshot();

  const [prevOpen, setPrevOpen] = useState(false);
  if (open !== prevOpen) {
    setPrevOpen(open);
    if (open) {
      setStep("select");
      setFilePath(null);
      setMetadata(null);
      setTruncate(false);
      reset();
    }
  }

  const [prevProgress, setPrevProgress] = useState(progress);
  if (progress !== prevProgress) {
    setPrevProgress(progress);
    if (progress?.stage === "done" || progress?.stage === "error") {
      setStep("done");
    }
  }

  const handleSelectFile = async () => {
    const selected = await openFile({
      filters: [{ name: "JSON Snapshot", extensions: ["json"] }],
      multiple: false,
    });
    if (!selected) return;
    const path = typeof selected === "string" ? selected : (selected as { path: string }).path;

    setFilePath(path);
    readMeta.mutate(path, {
      onSuccess: (meta) => {
        setMetadata(meta);
        setStep("confirm");
      },
      onError: (err) => toast.error("Invalid snapshot file", { description: String(err) }),
    });
  };

  const handleRestore = () => {
    if (!filePath) return;
    setStep("progress");

    const snapshotId = crypto.randomUUID();
    restore({
      params: {
        connection_id: connectionId,
        file_path: filePath,
        truncate_before_restore: truncate,
      },
      snapshotId,
    });
  };

  function formatBytes(bytes: number): string {
    if (bytes < 1024) return `${bytes} B`;
    if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed(0)} KB`;
    return `${(bytes / (1024 * 1024)).toFixed(1)} MB`;
  }

  return (
    <Dialog open={open} onOpenChange={onOpenChange}>
      <DialogContent className="sm:max-w-[520px]">
        <DialogHeader>
          <DialogTitle className="flex items-center gap-2">
            <Upload className="h-5 w-5" />
            Restore Snapshot
          </DialogTitle>
        </DialogHeader>

        {step === "select" && (
          <>
            <div className="py-8 flex flex-col items-center gap-3">
              <FileJson className="h-12 w-12 text-muted-foreground" />
              <p className="text-sm text-muted-foreground">Select a snapshot file to restore</p>
              <Button onClick={handleSelectFile} disabled={readMeta.isPending}>
                {readMeta.isPending ? (
                  <><Loader2 className="h-4 w-4 animate-spin mr-1" />Reading...</>
                ) : (
                  "Choose File"
                )}
              </Button>
            </div>
            <DialogFooter>
              <Button variant="outline" onClick={() => onOpenChange(false)}>Cancel</Button>
            </DialogFooter>
          </>
        )}

        {step === "confirm" && metadata && (
          <>
            <div className="space-y-3 py-2">
              <div className="rounded-md border p-3 space-y-2 text-sm">
                <div className="flex justify-between">
                  <span className="text-muted-foreground">Name</span>
                  <span className="font-medium">{metadata.name}</span>
                </div>
                <div className="flex justify-between">
                  <span className="text-muted-foreground">Created</span>
                  <span>{new Date(metadata.created_at).toLocaleString()}</span>
                </div>
                <div className="flex justify-between">
                  <span className="text-muted-foreground">Tables</span>
                  <span>{metadata.tables.length}</span>
                </div>
                <div className="flex justify-between">
                  <span className="text-muted-foreground">Total Rows</span>
                  <span>{metadata.total_rows.toLocaleString()}</span>
                </div>
                <div className="flex justify-between">
                  <span className="text-muted-foreground">File Size</span>
                  <span>{formatBytes(metadata.file_size_bytes)}</span>
                </div>
              </div>

              <div className="space-y-1">
                <p className="text-xs font-medium text-muted-foreground">Tables included:</p>
                <div className="flex flex-wrap gap-1">
                  {metadata.tables.map((t) => (
                    <Badge key={`${t.schema}.${t.table}`} variant="secondary" className="text-[10px]">
                      {t.schema}.{t.table} ({t.row_count})
                    </Badge>
                  ))}
                </div>
              </div>

              <div className="flex items-start gap-2 rounded-md border border-yellow-500/50 bg-yellow-500/10 p-3">
                <AlertTriangle className="h-4 w-4 text-yellow-600 shrink-0 mt-0.5" />
                <div className="space-y-2">
                  <label className="flex items-center gap-2 text-sm cursor-pointer">
                    <input
                      type="checkbox"
                      checked={truncate}
                      onChange={(e) => setTruncate(e.target.checked)}
                      className="rounded"
                    />
                    Truncate existing data before restore
                  </label>
                  {truncate && (
                    <p className="text-xs text-yellow-700 dark:text-yellow-400">
                      This will DELETE all existing data in the affected tables before restoring.
                    </p>
                  )}
                </div>
              </div>
            </div>

            <DialogFooter>
              <Button variant="outline" onClick={() => setStep("select")}>Back</Button>
              <Button onClick={handleRestore}>Restore</Button>
            </DialogFooter>
          </>
        )}

        {step === "progress" && (
          <div className="py-4 space-y-4">
            <div className="space-y-2">
              <div className="flex items-center justify-between text-sm">
                <span>{progress?.message || "Starting..."}</span>
                <span className="text-muted-foreground">{progress?.percent ?? 0}%</span>
              </div>
              <div className="h-2 rounded-full bg-secondary overflow-hidden">
                <div
                  className="h-full rounded-full bg-primary transition-all duration-300"
                  style={{ width: `${progress?.percent ?? 0}%` }}
                />
              </div>
            </div>
            {isRestoring && (
              <div className="flex items-center gap-2 text-sm text-muted-foreground">
                <Loader2 className="h-4 w-4 animate-spin" />
                {progress?.detail || progress?.stage || "Restoring..."}
              </div>
            )}
          </div>
        )}

        {step === "done" && (
          <div className="py-4 space-y-4">
            {error ? (
              <div className="flex items-start gap-3 rounded-md border border-destructive/50 bg-destructive/10 p-4">
                <XCircle className="h-5 w-5 text-destructive shrink-0 mt-0.5" />
                <div className="space-y-1">
                  <p className="text-sm font-medium text-destructive">Restore Failed</p>
                  <p className="text-xs text-muted-foreground">{error}</p>
                </div>
              </div>
            ) : (
              <div className="flex items-start gap-3 rounded-md border border-green-500/50 bg-green-500/10 p-4">
                <CheckCircle2 className="h-5 w-5 text-green-500 shrink-0 mt-0.5" />
                <div className="space-y-1">
                  <p className="text-sm font-medium">Restore Completed</p>
                  <p className="text-xs text-muted-foreground">
                    {rowsRestored?.toLocaleString()} rows restored successfully.
                  </p>
                </div>
              </div>
            )}

            <DialogFooter>
              <Button variant="outline" onClick={() => onOpenChange(false)}>Close</Button>
              {error && <Button onClick={() => setStep("confirm")}>Retry</Button>}
            </DialogFooter>
          </div>
        )}
      </DialogContent>
    </Dialog>
  );
}