import { Box, CircularProgress, Dialog, DialogTitle, Step, StepLabel, Stepper } from "@mui/material";
import { ThemeProvider, useTheme } from "@mui/material/styles";
import { useStore } from "@tanstack/react-form";
import { useSnackbar } from "notistack";
import { useCallback, useEffect, useMemo, useState } from "react";
import { useNavigate } from "react-router-dom";

import { EvalsStep } from "./components/evals-step";
import { InfoStep } from "./components/info-step";
import { PromptStep } from "./components/prompt-step";
import { createExperimentModalFormOpts, CreateExperimentModalFormValues } from "./form";
import { createExperimentDropdownTheme } from "./tourDropdownTheme";
import { DeepPartial, deepMerge, formDataToRequest, templateToFormData } from "./utils";

import { ConfirmationModal } from "@/components/common/ConfirmationModal";
import { useAppForm } from "@/components/traces/components/filtering/hooks/form";
import { TOUR_IDS, tourDataAttr } from "@/features/task-tour/selectors";
import { dispatchTourEvent, refreshTaskTourTarget, TASK_TOUR_EVENTS } from "@/features/task-tour/tourActions";
import { useCreateExperiment, usePromptExperiment } from "@/hooks/usePromptExperiments";
import { useTask } from "@/hooks/useTask";
import { PromptExperimentDetail } from "@/lib/api-client/api-client";

type Props = {
  templateId?: string;
  initialData?: DeepPartial<CreateExperimentModalFormValues>;
  open: boolean;
  onClose: () => void;
};

export const CreateExperimentModal = ({ templateId, initialData, open, onClose }: Props) => {
  const { experiment, isLoading } = usePromptExperiment(templateId);

  // Lift this modal's dropdown popups above the task tour overlay so the guided
  // "Create Experiment" step can open them without the tour backdrop/blocker
  // clipping them (UP-4482). See createExperimentDropdownTheme for details.
  const parentTheme = useTheme();
  const dropdownTheme = useMemo(() => createExperimentDropdownTheme(parentTheme), [parentTheme]);

  return (
    <ThemeProvider theme={dropdownTheme}>
      <Dialog
        open={open}
        maxWidth="md"
        fullWidth
        aria-labelledby="create-experiment-dialog-title"
        slotProps={{ paper: { ...tourDataAttr(TOUR_IDS.createExperimentModal) } }}
      >
        <DialogTitle id="create-experiment-dialog-title">Create Experiment</DialogTitle>
        {isLoading ? (
          <Box sx={{ display: "flex", justifyContent: "center", alignItems: "center", height: "100%", pb: 4 }}>
            <CircularProgress />
          </Box>
        ) : (
          <CreateExperimentModalInner template={experiment} initialData={initialData} onClose={onClose} />
        )}
      </Dialog>
    </ThemeProvider>
  );
};

const CreateExperimentModalInner = ({
  template,
  initialData,
  onClose,
}: {
  template?: PromptExperimentDetail;
  initialData?: DeepPartial<CreateExperimentModalFormValues>;
  onClose: () => void;
}) => {
  const { task } = useTask()!;
  const { enqueueSnackbar } = useSnackbar();
  const navigate = useNavigate();

  const [showDiscardConfirm, setShowDiscardConfirm] = useState(false);

  const createExperiment = useCreateExperiment(task!.id, {
    onSuccess: (data) => {
      dispatchTourEvent(TASK_TOUR_EVENTS.createExperimentCreated);
      enqueueSnackbar(`Experiment "${data.name}" created successfully!`, { variant: "success" });
      navigate(`/tasks/${task!.id}/prompt-experiments/${data.id}`);
      onClose();
    },
  });

  const defaultValues = (() => {
    const base = template ? templateToFormData(template) : createExperimentModalFormOpts.defaultValues;
    return initialData ? deepMerge(base, initialData) : base;
  })();

  const form = useAppForm({
    ...createExperimentModalFormOpts,
    defaultValues,
    onSubmit: async ({ value, formApi }) => {
      if (value.section === "info") {
        dispatchTourEvent(TASK_TOUR_EVENTS.createExperimentInfoCompleted);
        return formApi.setFieldValue("section", "prompts");
      }
      if (value.section === "prompts") {
        dispatchTourEvent(TASK_TOUR_EVENTS.createExperimentPromptMappingsCompleted);
        // Skip the evals step entirely when no evaluators were selected — there
        // would be nothing to configure on that screen.
        if (value.info.evaluators.length === 0) {
          const request = { ...formDataToRequest(value), eval_list: [] };
          await createExperiment.mutateAsync(request);
          formApi.reset();
          return;
        }
        return formApi.setFieldValue("section", "evals");
      }

      const request = formDataToRequest(value);

      await createExperiment.mutateAsync(request);

      formApi.reset();
    },
  });

  const handleCancel = useCallback(() => {
    if (form.state.isDirty) {
      setShowDiscardConfirm(true);
    } else {
      onClose();
    }
  }, [form.state.isDirty, onClose]);

  const section = useStore(form.store, (state) => state.values.section);
  const hasEvaluators = useStore(form.store, (state) => state.values.info.evaluators.length > 0);

  useEffect(() => {
    const frame = window.requestAnimationFrame(() => refreshTaskTourTarget());
    return () => window.cancelAnimationFrame(frame);
  }, [section]);

  const step = {
    info: 0,
    prompts: 1,
    evals: 2,
  }[section];

  return (
    <>
      <form
        onSubmit={(e) => {
          e.preventDefault();
          e.stopPropagation();
          form.handleSubmit();
        }}
        className="contents"
      >
        <Stepper activeStep={step} sx={{ p: 2 }}>
          <Step>
            <StepLabel>Experiment Info</StepLabel>
          </Step>
          <Step>
            <StepLabel>Configure Prompts</StepLabel>
          </Step>
          {hasEvaluators && (
            <Step>
              <StepLabel>Configure Evals</StepLabel>
            </Step>
          )}
        </Stepper>
        {section === "info" && (
          <Box data-tour-id={TOUR_IDS.createExperimentInfoStep}>
            <InfoStep form={form} onCancel={handleCancel} />
          </Box>
        )}
        {section === "prompts" && (
          <Box data-tour-id={TOUR_IDS.createExperimentPromptMappingsStep}>
            <PromptStep form={form} onCancel={handleCancel} />
          </Box>
        )}
        {section === "evals" && (
          <Box data-tour-id={TOUR_IDS.createExperimentEvalMappingsStep}>
            <EvalsStep form={form} onCancel={handleCancel} />
          </Box>
        )}
      </form>
      <ConfirmationModal
        open={showDiscardConfirm}
        onClose={() => setShowDiscardConfirm(false)}
        onConfirm={onClose}
        title="Discard changes?"
        message="You have unsaved changes. Are you sure you want to close this form? Your changes will be lost."
        confirmText="Discard"
        cancelText="Keep editing"
      />
    </>
  );
};
