tusk/src/components/data-generator/GenerateDataDialog.tsx

import { useState } from "react";
import {
  Dialog,
  DialogContent,
  DialogHeader,
  DialogTitle,
  DialogFooter,
} from "@/components/ui/dialog";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Badge } from "@/components/ui/badge";
import { useDataGenerator } from "@/hooks/use-data-generator";
import { toast } from "sonner";
import {
  Loader2,
  CheckCircle2,
  XCircle,
  Wand2,
  Table2,
} from "lucide-react";

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

type Step = "config" | "preview" | "done";

export function GenerateDataDialog({
  open,
  onOpenChange,
  connectionId,
  schema,
  table,
}: Props) {
  const [step, setStep] = useState<Step>("config");
  const [rowCount, setRowCount] = useState(10);
  const [includeRelated, setIncludeRelated] = useState(true);
  const [customInstructions, setCustomInstructions] = useState("");

  const {
    generatePreview,
    preview,
    isGenerating,
    generateError,
    insertData,
    insertedRows,
    isInserting,
    insertError,
    progress,
    reset,
  } = useDataGenerator();

  const [prevOpen, setPrevOpen] = useState(false);
  if (open !== prevOpen) {
    setPrevOpen(open);
    if (open) {
      setStep("config");
      setRowCount(10);
      setIncludeRelated(true);
      setCustomInstructions("");
      reset();
    }
  }

  const handleGenerate = () => {
    const genId = crypto.randomUUID();
    generatePreview(
      {
        params: {
          connection_id: connectionId,
          schema,
          table,
          row_count: rowCount,
          include_related: includeRelated,
          custom_instructions: customInstructions || undefined,
        },
        genId,
      },
      {
        onSuccess: () => setStep("preview"),
        onError: (err) => toast.error("Generation failed", { description: String(err) }),
      }
    );
  };

  const handleInsert = () => {
    if (!preview) return;
    insertData(
      { connectionId, preview },
      {
        onSuccess: (rows) => {
          setStep("done");
          toast.success(`Inserted ${rows} rows`);
        },
        onError: (err) => toast.error("Insert failed", { description: String(err) }),
      }
    );
  };

  return (
    <Dialog open={open} onOpenChange={onOpenChange}>
      <DialogContent className="sm:max-w-[600px] max-h-[80vh] overflow-y-auto">
        <DialogHeader>
          <DialogTitle className="flex items-center gap-2">
            <Wand2 className="h-5 w-5" />
            Generate Test Data
          </DialogTitle>
        </DialogHeader>

        {step === "config" && (
          <>
            <div className="grid gap-3 py-2">
              <div className="grid grid-cols-4 items-center gap-3">
                <label className="text-right text-sm text-muted-foreground">Table</label>
                <div className="col-span-3">
                  <Badge variant="secondary">{schema}.{table}</Badge>
                </div>
              </div>

              <div className="grid grid-cols-4 items-center gap-3">
                <label className="text-right text-sm text-muted-foreground">Row Count</label>
                <Input
                  className="col-span-3"
                  type="number"
                  value={rowCount}
                  onChange={(e) => setRowCount(Math.min(1000, Math.max(1, parseInt(e.target.value) || 1)))}
                  min={1}
                  max={1000}
                />
              </div>

              <div className="grid grid-cols-4 items-center gap-3">
                <label className="text-right text-sm text-muted-foreground">Related Tables</label>
                <div className="col-span-3 flex items-center gap-2">
                  <input
                    type="checkbox"
                    checked={includeRelated}
                    onChange={(e) => setIncludeRelated(e.target.checked)}
                    className="rounded"
                  />
                  <span className="text-sm text-muted-foreground">
                    Include parent tables (via foreign keys)
                  </span>
                </div>
              </div>

              <div className="grid grid-cols-4 items-start gap-3">
                <label className="text-right text-sm text-muted-foreground pt-2">Instructions</label>
                <Input
                  className="col-span-3"
                  placeholder="Optional: specific data requirements..."
                  value={customInstructions}
                  onChange={(e) => setCustomInstructions(e.target.value)}
                />
              </div>
            </div>

            {isGenerating && progress && (
              <div className="space-y-2">
                <div className="flex items-center justify-between text-sm">
                  <span>{progress.message}</span>
                  <span className="text-muted-foreground">{progress.percent}%</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}%` }}
                  />
                </div>
              </div>
            )}

            <DialogFooter>
              <Button variant="outline" onClick={() => onOpenChange(false)}>Cancel</Button>
              <Button onClick={handleGenerate} disabled={isGenerating}>
                {isGenerating ? (
                  <><Loader2 className="h-4 w-4 animate-spin mr-1" />Generating...</>
                ) : (
                  "Generate Preview"
                )}
              </Button>
            </DialogFooter>
          </>
        )}

        {step === "preview" && preview && (
          <>
            <div className="space-y-3">
              <div className="flex items-center gap-2 text-sm">
                <span className="text-muted-foreground">Preview:</span>
                <Badge variant="secondary">{preview.total_rows} rows across {preview.tables.length} tables</Badge>
              </div>

              {preview.tables.map((tbl) => (
                <div key={`${tbl.schema}.${tbl.table}`} className="rounded-md border">
                  <div className="flex items-center gap-2 px-3 py-2 bg-muted/50 text-sm font-medium border-b">
                    <Table2 className="h-3.5 w-3.5" />
                    {tbl.schema}.{tbl.table}
                    <Badge variant="secondary" className="ml-auto text-[10px]">{tbl.row_count} rows</Badge>
                  </div>
                  <div className="overflow-x-auto max-h-48">
                    <table className="w-full text-xs">
                      <thead>
                        <tr className="border-b">
                          {tbl.columns.map((col) => (
                            <th key={col} className="px-2 py-1 text-left font-medium text-muted-foreground whitespace-nowrap">
                              {col}
                            </th>
                          ))}
                        </tr>
                      </thead>
                      <tbody>
                        {tbl.rows.slice(0, 5).map((row, i) => (
                          <tr key={i} className="border-b last:border-0">
                            {(row as unknown[]).map((val, j) => (
                              <td key={j} className="px-2 py-1 font-mono whitespace-nowrap">
                                {val === null ? (
                                  <span className="text-muted-foreground">NULL</span>
                                ) : (
                                  String(val).substring(0, 50)
                                )}
                              </td>
                            ))}
                          </tr>
                        ))}
                        {tbl.rows.length > 5 && (
                          <tr>
                            <td colSpan={tbl.columns.length} className="px-2 py-1 text-center text-muted-foreground">
                              ...and {tbl.rows.length - 5} more rows
                            </td>
                          </tr>
                        )}
                      </tbody>
                    </table>
                  </div>
                </div>
              ))}
            </div>

            <DialogFooter>
              <Button variant="outline" onClick={() => setStep("config")}>Back</Button>
              <Button onClick={handleInsert} disabled={isInserting}>
                {isInserting ? (
                  <><Loader2 className="h-4 w-4 animate-spin mr-1" />Inserting...</>
                ) : (
                  `Insert ${preview.total_rows} Rows`
                )}
              </Button>
            </DialogFooter>
          </>
        )}

        {step === "done" && (
          <div className="py-4 space-y-4">
            {insertError ? (
              <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">Insert Failed</p>
                  <p className="text-xs text-muted-foreground">{insertError}</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">Data Generated Successfully</p>
                  <p className="text-xs text-muted-foreground">
                    {insertedRows} rows inserted across {preview?.tables.length ?? 0} tables.
                  </p>
                </div>
              </div>
            )}

            <DialogFooter>
              <Button variant="outline" onClick={() => onOpenChange(false)}>Close</Button>
              {insertError && (
                <Button onClick={() => setStep("preview")}>Retry</Button>
              )}
            </DialogFooter>
          </div>
        )}

        {generateError && step === "config" && (
          <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" />
            <p className="text-xs text-muted-foreground">{generateError}</p>
          </div>
        )}
      </DialogContent>
    </Dialog>
  );
}