import { AutoAwesome, Close } from "@mui/icons-material";
import { Box, Dialog, DialogContent, DialogTitle, IconButton, Typography } from "@mui/material";
import React, { useCallback, useEffect, useRef, useState } from "react";

import { SyntheticDataCanvas } from "./SyntheticDataCanvas";
import { SyntheticDataConfigForm } from "./SyntheticDataConfigForm";
import type { GenerationConfig } from "./types";

import { ConfirmationModal } from "@/components/common/ConfirmationModal";
import { TOUR_IDS, tourDataAttr } from "@/features/task-tour/selectors";
import { dispatchTourEvent, refreshTaskTourTarget, TASK_TOUR_EVENTS } from "@/features/task-tour/tourActions";
import { useSyntheticDataSession } from "@/hooks/datasets/useSyntheticDataSession";
import type { DatasetVersionRowResponse } from "@/lib/api-client/api-client";

interface SyntheticDataModalProps {
  open: boolean;
  onClose: () => void;
  columns: string[];
  existingRowsSample: DatasetVersionRowResponse[];
  datasetId: string;
  versionNumber: number;
  onAcceptRows: (rows: { data: { column_name: string; column_value: string }[] }[]) => void;
}

type ModalPhase = "configure" | "canvas";

export const SyntheticDataModal: React.FC<SyntheticDataModalProps> = ({
  open,
  onClose,
  columns,
  existingRowsSample,
  datasetId,
  versionNumber,
  onAcceptRows,
}) => {
  const [phase, setPhase] = useState<ModalPhase>("configure");
  const [config, setConfig] = useState<GenerationConfig | null>(null);
  const [showDiscardConfirm, setShowDiscardConfirm] = useState(false);
  const completionDispatchedRef = useRef(false);

  const session = useSyntheticDataSession(datasetId, versionNumber, columns);

  useEffect(() => {
    if (!open) return;
    completionDispatchedRef.current = false;
    const frame = window.requestAnimationFrame(() => refreshTaskTourTarget());
    return () => window.cancelAnimationFrame(frame);
  }, [open]);

  const completeSyntheticTourStep = useCallback(() => {
    if (completionDispatchedRef.current) return;
    completionDispatchedRef.current = true;
    dispatchTourEvent(TASK_TOUR_EVENTS.syntheticDataFinished);
  }, []);

  const handleStartGeneration = useCallback(
    async (generationConfig: GenerationConfig) => {
      setConfig(generationConfig);
      await session.startGeneration(generationConfig, existingRowsSample);
      completeSyntheticTourStep();
      setPhase("canvas");
    },
    [session, existingRowsSample, completeSyntheticTourStep]
  );

  const handleSendMessage = useCallback(
    async (message: string) => {
      if (config) {
        await session.sendMessage(message, config);
      }
    },
    [session, config]
  );

  const handleClose = useCallback(() => {
    completeSyntheticTourStep();
    setPhase("configure");
    setConfig(null);
    session.reset();
    onClose();
  }, [session, onClose, completeSyntheticTourStep]);

  const handleAcceptRows = useCallback(
    (selectedIds: Set<string>) => {
      // Convert selected synthetic rows to the format expected by the dataset
      const rowsToAdd = session.rows
        .filter((row) => selectedIds.has(row.id))
        .map((row) => ({
          data: Object.entries(row.data).map(([column_name, column_value]) => ({
            column_name,
            column_value,
          })),
        }));
      onAcceptRows(rowsToAdd);
      handleClose();
    },
    [session.rows, onAcceptRows, handleClose]
  );

  const handleAttemptClose = useCallback(() => {
    if (phase === "canvas" && session.rows.length > 0) {
      setShowDiscardConfirm(true);
    } else {
      handleClose();
    }
  }, [phase, session.rows.length, handleClose]);

  const handleBack = useCallback(() => {
    if (session.rows.length > 0) {
      setShowDiscardConfirm(true);
    } else {
      setPhase("configure");
      session.reset();
    }
  }, [session]);

  return (
    <Dialog
      open={open}
      onClose={handleAttemptClose}
      maxWidth={phase === "canvas" ? "xl" : "md"}
      fullWidth
      slotProps={{
        paper: {
          ...tourDataAttr(TOUR_IDS.datasetGenerateSyntheticModal),
          sx: {
            height: phase === "canvas" ? "90vh" : "auto",
            maxHeight: phase === "canvas" ? "90vh" : "auto",
          },
        },
      }}
    >
      <DialogTitle
        sx={{
          display: "flex",
          alignItems: "center",
          justifyContent: "space-between",
          pb: 1,
        }}
      >
        <Box sx={{ display: "flex", alignItems: "center", gap: 1 }}>
          <AutoAwesome color="primary" />
          <Typography variant="h6">{phase === "configure" ? "Generate Synthetic Data" : "Synthetic Data Generation"}</Typography>
        </Box>
        <IconButton onClick={handleAttemptClose} size="small">
          <Close />
        </IconButton>
      </DialogTitle>

      <DialogContent
        sx={{
          display: "flex",
          flexDirection: "column",
          height: phase === "canvas" ? "calc(100% - 64px)" : "auto",
          p: phase === "canvas" ? 0 : 3,
        }}
      >
        {phase === "configure" ? (
          <SyntheticDataConfigForm
            columns={columns}
            existingRowsSample={existingRowsSample}
            onSubmit={handleStartGeneration}
            onCancel={handleClose}
            isLoading={session.isLoading}
          />
        ) : (
          <SyntheticDataCanvas
            rows={session.rows}
            conversation={session.conversation}
            columns={columns}
            isLoading={session.isLoading}
            error={session.error}
            onSendMessage={handleSendMessage}
            onUpdateRow={session.updateRow}
            onAddRow={session.addRow}
            onDeleteRows={session.deleteRows}
            onToggleLock={session.toggleLock}
            onBack={handleBack}
            onAccept={handleAcceptRows}
          />
        )}
      </DialogContent>

      <ConfirmationModal
        open={showDiscardConfirm}
        onClose={() => setShowDiscardConfirm(false)}
        onConfirm={handleClose}
        title="Discard Generated Data?"
        message="You have generated data that hasn't been accepted. Closing will discard all generated rows and conversation history."
        confirmText="Discard"
        cancelText="Keep Editing"
      />
    </Dialog>
  );
};
