import AddIcon from "@mui/icons-material/Add";
import CloseIcon from "@mui/icons-material/Close";
import InfoOutlinedIcon from "@mui/icons-material/InfoOutlined";
import {
  Dialog,
  DialogTitle,
  DialogContent,
  DialogActions,
  Button,
  TextField,
  Box,
  CircularProgress,
  Autocomplete,
  Chip,
  FormControl,
  InputLabel,
  Select,
  MenuItem,
  Typography,
  Stepper,
  Step,
  StepLabel,
  Tooltip,
  ToggleButtonGroup,
  ToggleButton,
  Snackbar,
  Alert,
} from "@mui/material";
import { alpha } from "@mui/material/styles";
import React, { useState, useEffect } from "react";
import { useParams, useNavigate } from "react-router-dom";

import { useApi } from "@/hooks/useApi";
import useSnackbar from "@/hooks/useSnackbar";
import type {
  LLMGetAllMetadataResponse,
  AgenticPromptVersionResponse,
  DatasetResponse,
  DatasetVersionMetadataResponse,
  LLMVersionResponse,
  NewDatasetVersionRowColumnItemRequest,
  PromptExperimentDetail,
} from "@/lib/api-client/api-client";

interface CreateExperimentModalProps {
  open: boolean;
  onClose: () => void;
  onSubmit: (data: ExperimentFormData) => Promise<{ id: string }>;
  initialData?: Partial<PromptExperimentDetail>;
  isLoadingInitialData?: boolean;
  disableNavigation?: boolean; // Don't navigate after creation (for notebook config mode)
}

export interface PromptVersionSelection {
  promptName: string;
  version: number;
}

export interface EvaluatorSelection {
  name: string;
  version: number;
}

export type VariableSourceType = "dataset_column" | "experiment_output";

export interface VariableMapping {
  variableName: string;
  sourceType: VariableSourceType;
  datasetColumn?: string; // For dataset_column source
  jsonPath?: string; // For experiment_output source
}

export interface PromptVariableMappings {
  [variableName: string]: string; // variable name -> dataset column name
}

export interface EvalVariableMappings {
  evalName: string;
  evalVersion: number;
  mappings: {
    [variableName: string]: {
      sourceType: VariableSourceType;
      datasetColumn?: string;
      jsonPath?: string;
    };
  };
}

export type DatasetRowFilter = NewDatasetVersionRowColumnItemRequest;

export interface ExperimentFormData {
  name: string;
  description: string;
  promptVersions: PromptVersionSelection[];
  datasetId: string;
  datasetName: string;
  datasetVersion: number | "";
  evaluators: EvaluatorSelection[];
  promptVariableMappings?: PromptVariableMappings;
  evalVariableMappings?: EvalVariableMappings[];
  datasetRowFilter?: DatasetRowFilter[];
}

export const CreateExperimentModal: React.FC<CreateExperimentModalProps> = ({
  open,
  onClose,
  onSubmit,
  initialData,
  isLoadingInitialData = false,
  disableNavigation = false,
}) => {
  const { id: taskId } = useParams<{ id: string }>();
  const navigate = useNavigate();
  const api = useApi();
  const { showSnackbar, snackbarProps, alertProps } = useSnackbar();

  // Step management
  const [currentStep, setCurrentStep] = useState(0);
  const [completedSteps, setCompletedSteps] = useState<Set<number>>(new Set());
  const [currentEvalIndex, setCurrentEvalIndex] = useState(0); // Track which eval is being configured in step 2
  const [isInitializing, setIsInitializing] = useState(false); // Track internal data loading when cloning

  // Form state
  const [formData, setFormData] = useState<ExperimentFormData>({
    name: "",
    description: "",
    promptVersions: [],
    datasetId: "",
    datasetName: "",
    datasetVersion: "",
    evaluators: [],
    promptVariableMappings: {},
    evalVariableMappings: [],
    datasetRowFilter: [],
  });

  // Prompts state
  const [prompts, setPrompts] = useState<LLMGetAllMetadataResponse[]>([]);
  const [selectedPromptName, setSelectedPromptName] = useState<string>("");
  const [promptVersions, setPromptVersions] = useState<AgenticPromptVersionResponse[]>([]);
  const [visibleOlderVersions, setVisibleOlderVersions] = useState<number[]>([]);
  const [loadingPrompts, setLoadingPrompts] = useState(false);
  const [loadingPromptVersions, setLoadingPromptVersions] = useState(false);

  // Datasets state
  const [datasets, setDatasets] = useState<DatasetResponse[]>([]);
  const [datasetVersions, setDatasetVersions] = useState<DatasetVersionMetadataResponse[]>([]);
  const [datasetColumns, setDatasetColumns] = useState<string[]>([]);
  const [loadingDatasets, setLoadingDatasets] = useState(false);
  const [loadingDatasetVersions, setLoadingDatasetVersions] = useState(false);

  // Prompt and eval variable details
  const [promptVariables, setPromptVariables] = useState<string[]>([]);
  const [loadingPromptDetails, setLoadingPromptDetails] = useState(false);
  const [evalVariables, setEvalVariables] = useState<Record<string, { name: string; version: number; variables: string[] }>>({});
  const [loadingEvalDetails, setLoadingEvalDetails] = useState(false);

  // Evaluator instructions modal state
  const [instructionsModalOpen, setInstructionsModalOpen] = useState(false);
  const [selectedEvalInstructions, setSelectedEvalInstructions] = useState<{ name: string; version: number; instructions: string } | null>(null);
  const [loadingInstructions, setLoadingInstructions] = useState(false);

  // Evaluators state
  const [evaluators, setEvaluators] = useState<LLMGetAllMetadataResponse[]>([]);
  const [evaluatorVersions, setEvaluatorVersions] = useState<Record<string, LLMVersionResponse[]>>({});
  const [currentEvaluatorName, setCurrentEvaluatorName] = useState<string>("");
  const [currentEvaluatorVersion, setCurrentEvaluatorVersion] = useState<number | "">("");

  const [errors, setErrors] = useState<Partial<Record<keyof ExperimentFormData | "general", string>>>({});
  const [isSubmitting, setIsSubmitting] = useState(false);

  // Load prompts on mount
  useEffect(() => {
    if (open && taskId && api) {
      loadPrompts();
      loadDatasets();
      loadEvaluators();
    }
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [open, taskId, api]);

  // Load prompt versions when a prompt is selected
  useEffect(() => {
    if (selectedPromptName && taskId && api) {
      loadPromptVersions(selectedPromptName);
    }
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [selectedPromptName, taskId, api]);

  // Initialize form from existing experiment data
  useEffect(() => {
    const initializeFromExistingExperiment = async () => {
      if (!initialData || !open || isLoadingInitialData || !taskId || !api) {
        // If we don't have initialData, we're not initializing
        if (!initialData && open) {
          setIsInitializing(false);
        }
        return;
      }

      // Set initializing to true at the start of the async function
      setIsInitializing(true);

      try {
        // Check if we're initializing from a full experiment or just a dataset
        const isFullExperiment = initialData.prompt_configs && initialData.eval_list;

        // Transform the initial data to form data format
        const promptVersions: PromptVersionSelection[] =
          isFullExperiment && initialData.prompt_configs
            ? initialData.prompt_configs
                .filter((pc): pc is { type: "saved" } & { name: string; version: number } => pc.type === "saved")
                .map((pc) => ({
                  promptName: pc.name,
                  version: pc.version,
                }))
            : [];

        const evaluators: EvaluatorSelection[] =
          isFullExperiment && initialData.eval_list
            ? initialData.eval_list.map((e) => ({
                name: e.name,
                version: e.version,
              }))
            : [];

        // Transform prompt variable mappings
        const promptVariableMappings: PromptVariableMappings = {};
        if (isFullExperiment && initialData.prompt_variable_mapping) {
          initialData.prompt_variable_mapping.forEach((mapping) => {
            if (mapping.source.type === "dataset_column") {
              promptVariableMappings[mapping.variable_name] = mapping.source.dataset_column.name;
            }
          });
        }

        // Transform eval variable mappings
        const evalVariableMappings: EvalVariableMappings[] =
          isFullExperiment && initialData.eval_list
            ? initialData.eval_list.map((evalConfig) => {
                const mappings: EvalVariableMappings["mappings"] = {};
                evalConfig.variable_mapping.forEach((mapping) => {
                  if (mapping.source.type === "dataset_column") {
                    mappings[mapping.variable_name] = {
                      sourceType: "dataset_column",
                      datasetColumn: mapping.source.dataset_column.name,
                    };
                  } else if (mapping.source.type === "experiment_output") {
                    mappings[mapping.variable_name] = {
                      sourceType: "experiment_output",
                      jsonPath: mapping.source.experiment_output.json_path || "",
                    };
                  }
                });
                return {
                  evalName: evalConfig.name,
                  evalVersion: evalConfig.version,
                  mappings,
                };
              })
            : [];

        // Set the selected prompt name first (from first saved prompt config)
        const firstSavedPrompt = initialData.prompt_configs?.find(
          (pc): pc is { type: "saved" } & { name: string; version: number } => pc.type === "saved"
        );
        if (isFullExperiment && firstSavedPrompt) {
          setSelectedPromptName(firstSavedPrompt.name);
        }

        // Load all necessary data in parallel
        // Pass the desired version to preserve the original dataset version
        const loadTasks = [loadDatasetVersions(initialData.dataset_ref!.id, initialData.dataset_ref!.version)];

        if (isFullExperiment && firstSavedPrompt) {
          loadTasks.push(loadPromptVersions(firstSavedPrompt.name));
          loadTasks.push(...evaluators.map((evaluator) => loadEvaluatorVersions(evaluator.name)));
        }

        await Promise.all(loadTasks);

        // Transform dataset row filter from initial data
        const datasetRowFilter: DatasetRowFilter[] = isFullExperiment && initialData.dataset_row_filter ? initialData.dataset_row_filter : [];

        // Now set the form data after all dropdowns are populated
        setFormData({
          name: isFullExperiment && initialData.name ? `${initialData.name} (Copy)` : "",
          description: isFullExperiment ? initialData.description || "" : "",
          promptVersions,
          datasetId: initialData.dataset_ref!.id,
          datasetName: initialData.dataset_ref!.name,
          datasetVersion: initialData.dataset_ref!.version,
          evaluators,
          promptVariableMappings,
          evalVariableMappings,
          datasetRowFilter,
        });

        // Clear the "add evaluator" form state that was set during loadEvaluatorVersions
        setCurrentEvaluatorName("");
        setCurrentEvaluatorVersion("");
      } catch (error) {
        console.error("Failed to initialize from existing experiment:", error);
      } finally {
        setIsInitializing(false);
      }
    };

    initializeFromExistingExperiment();
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [initialData, open, isLoadingInitialData]);

  const loadPrompts = async () => {
    if (!taskId || !api) return;
    try {
      setLoadingPrompts(true);
      const response = await api.api.getAllAgenticPromptsApiV1TasksTaskIdPromptsGet({
        taskId,
        page_size: 100,
      });
      setPrompts(response.data.llm_metadata);
    } catch (error) {
      console.error("Failed to load prompts:", error);
    } finally {
      setLoadingPrompts(false);
    }
  };

  const loadPromptVersions = async (promptName: string) => {
    if (!taskId || !api) return;
    try {
      setLoadingPromptVersions(true);
      const response = await api.api.getAllAgenticPromptVersionsApiV1TasksTaskIdPromptsPromptNameVersionsGet({
        taskId,
        promptName,
        page_size: 100,
      });
      const versions = response.data.versions.filter((v) => !v.deleted_at);
      setPromptVersions(versions);

      // Auto-select the version if there is only one available
      if (versions.length === 1) {
        setFormData((prev) => {
          const alreadySelected = prev.promptVersions.some((pv) => pv.promptName === promptName && pv.version === versions[0].version);
          if (alreadySelected) return prev;
          return {
            ...prev,
            promptVersions: [...prev.promptVersions, { promptName, version: versions[0].version }],
          };
        });
      }
    } catch (error) {
      console.error("Failed to load prompt versions:", error);
    } finally {
      setLoadingPromptVersions(false);
    }
  };

  const loadDatasets = async () => {
    if (!api || !taskId) return;
    try {
      setLoadingDatasets(true);
      const response = await api.api.getDatasetsApiV2TasksTaskIdDatasetsSearchGet({
        taskId,
        page_size: 100,
      });
      setDatasets(response.data.datasets);
    } catch (error) {
      console.error("Failed to load datasets:", error);
    } finally {
      setLoadingDatasets(false);
    }
  };

  const loadDatasetVersions = async (datasetId: string, desiredVersion?: number) => {
    if (!api) return;
    try {
      setLoadingDatasetVersions(true);
      const response = await api.api.getDatasetVersionsApiV2DatasetsDatasetIdVersionsGet({
        datasetId,
        page_size: 100,
      });
      const versions = response.data.versions;
      setDatasetVersions(versions);

      if (versions.length > 0) {
        // If a desired version is specified (during initialization), use it
        // Otherwise, default to the highest version number
        const targetVersion = desiredVersion ?? Math.max(...versions.map((v) => v.version_number));

        setFormData((prev) => ({
          ...prev,
          datasetVersion: targetVersion,
        }));

        // Load columns for the target version
        const versionData = versions.find((v) => v.version_number === targetVersion);
        if (versionData) {
          setDatasetColumns(versionData.column_names);
        }
      }
    } catch (error) {
      console.error("Failed to load dataset versions:", error);
    } finally {
      setLoadingDatasetVersions(false);
    }
  };

  const loadEvaluators = async () => {
    if (!taskId || !api) return;
    try {
      const response = await api.api.getAllLlmEvalsApiV1TasksTaskIdLlmEvalsGet({
        taskId,
        page_size: 100,
      });
      setEvaluators(response.data.llm_metadata);
    } catch (error) {
      console.error("Failed to load evaluators:", error);
    }
  };

  const loadEvaluatorVersions = async (evalName: string) => {
    if (!taskId || !api) return;
    try {
      const response = await api.api.getAllLlmEvalVersionsApiV1TasksTaskIdLlmEvalsEvalNameVersionsGet({
        taskId,
        evalName,
        page_size: 100,
      });
      const versions = response.data.versions.filter((v) => !v.deleted_at);
      setEvaluatorVersions((prev) => ({
        ...prev,
        [evalName]: versions,
      }));

      // Set to highest version number
      if (versions.length > 0) {
        const maxVersion = Math.max(...versions.map((v) => v.version));
        setCurrentEvaluatorVersion(maxVersion);
      }
    } catch (error) {
      console.error("Failed to load evaluator versions:", error);
    }
  };

  const loadEvalVariablesForEvaluator = async (evalName: string, evalVersion: number) => {
    if (!taskId || !api) return;
    try {
      setLoadingEvalDetails(true);
      const response = await api.api.getLlmEvalApiV1TasksTaskIdLlmEvalsEvalNameVersionsEvalVersionGet(evalName, String(evalVersion), taskId);
      if (response.data.variables) {
        setEvalVariables((prev) => ({
          ...prev,
          [`${evalName}-${evalVersion}`]: {
            name: evalName,
            version: evalVersion,
            variables: response.data.variables || [],
          },
        }));
      }
    } catch (error) {
      console.error("Failed to load eval variables:", error);
    } finally {
      setLoadingEvalDetails(false);
    }
  };

  const loadEvaluatorInstructions = async (evalName: string, evalVersion: number) => {
    if (!taskId || !api) return;
    try {
      setLoadingInstructions(true);
      const response = await api.api.getLlmEvalApiV1TasksTaskIdLlmEvalsEvalNameVersionsEvalVersionGet(evalName, String(evalVersion), taskId);
      setSelectedEvalInstructions({
        name: evalName,
        version: evalVersion,
        instructions: response.data.instructions || "",
      });
      setInstructionsModalOpen(true);
    } catch (error) {
      console.error("Failed to load evaluator instructions:", error);
    } finally {
      setLoadingInstructions(false);
    }
  };

  const handleAddPromptVersion = (version: number) => {
    if (!selectedPromptName) return;

    const existingIndex = formData.promptVersions.findIndex((pv) => pv.promptName === selectedPromptName && pv.version === version);

    if (existingIndex >= 0) {
      // Remove if already selected (toggle off)
      setFormData((prev) => ({
        ...prev,
        promptVersions: prev.promptVersions.filter((_, i) => i !== existingIndex),
      }));

      // Also remove from visible older versions if it's not in the top 5
      const top5Versions = promptVersions.slice(0, 5).map((v) => v.version);
      if (!top5Versions.includes(version)) {
        setVisibleOlderVersions((prev) => prev.filter((v) => v !== version));
      }
    } else {
      // Add if not selected (toggle on)
      setFormData((prev) => ({
        ...prev,
        promptVersions: [...prev.promptVersions, { promptName: selectedPromptName, version }],
      }));
    }
  };

  const handleAddEvaluator = () => {
    if (!currentEvaluatorName || !currentEvaluatorVersion) return;

    const alreadyAdded = formData.evaluators.some((e) => e.name === currentEvaluatorName && e.version === currentEvaluatorVersion);

    if (!alreadyAdded) {
      setFormData((prev) => ({
        ...prev,
        evaluators: [...prev.evaluators, { name: currentEvaluatorName, version: currentEvaluatorVersion as number }],
      }));

      // Clear the current selection
      setCurrentEvaluatorName("");
      setCurrentEvaluatorVersion("");
    }
  };

  const handleRemoveEvaluator = (index: number) => {
    setFormData((prev) => ({
      ...prev,
      evaluators: prev.evaluators.filter((_, i) => i !== index),
    }));
  };

  const validate = (): boolean => {
    const newErrors: Partial<Record<keyof ExperimentFormData | "general", string>> = {};

    if (!formData.name.trim()) {
      newErrors.name = "Experiment name is required";
    }

    if (formData.promptVersions.length === 0) {
      newErrors.promptVersions = "At least one prompt version is required";
    }

    if (!formData.datasetId) {
      newErrors.datasetId = "Dataset is required";
    }

    if (!formData.datasetVersion) {
      newErrors.datasetVersion = "Dataset version is required";
    }

    if (formData.evaluators.length === 0) {
      newErrors.evaluators = "At least one evaluator is required";
    }

    setErrors(newErrors);
    return Object.keys(newErrors).length === 0;
  };

  const handleSubmit = async () => {
    if (!validate()) return;

    try {
      setIsSubmitting(true);
      // Transform formData to match API expectations
      // The parent component will handle the actual API call
      const result = await onSubmit(formData);
      // Show success toast
      showSnackbar(`Experiment "${formData.name}" created successfully!`, "success");
      // Reset form on success
      setFormData({
        name: "",
        description: "",
        promptVersions: [],
        datasetId: "",
        datasetName: "",
        datasetVersion: "",
        evaluators: [],
        promptVariableMappings: {},
        evalVariableMappings: [],
        datasetRowFilter: [],
      });
      setSelectedPromptName("");
      setVisibleOlderVersions([]);
      setCurrentEvaluatorName("");
      setCurrentEvaluatorVersion("");
      setPromptVariables([]);
      setEvalVariables({});
      setDatasetColumns([]);
      setCurrentStep(0);
      setCurrentEvalIndex(0);
      setCompletedSteps(new Set());
      setErrors({});
      onClose();
      // Navigate to the experiment detail page (unless disabled for notebook config mode)
      if (!disableNavigation) {
        navigate(`/tasks/${taskId}/prompt-experiments/${result.id}`);
      }
    } catch (error) {
      console.error("Failed to create experiment:", error);
      setErrors({ general: "Failed to create experiment. Please try again." });
      showSnackbar("Failed to create experiment. Please try again.", "error");
    } finally {
      setIsSubmitting(false);
    }
  };

  const handleCancel = () => {
    // Reset form on cancel
    setFormData({
      name: "",
      description: "",
      promptVersions: [],
      datasetId: "",
      datasetName: "",
      datasetVersion: "",
      evaluators: [],
      promptVariableMappings: {},
      evalVariableMappings: [],
      datasetRowFilter: [],
    });
    setSelectedPromptName("");
    setVisibleOlderVersions([]);
    setCurrentEvaluatorName("");
    setCurrentEvaluatorVersion("");
    setPromptVariables([]);
    setEvalVariables({});
    setDatasetColumns([]);
    setCurrentStep(0);
    setCurrentEvalIndex(0);
    setCompletedSteps(new Set());
    setErrors({});
    setIsInitializing(false);
    onClose();
  };

  // Step navigation - always show 3 high-level steps
  const getStepLabels = () => {
    return ["Experiment Info", "Configure Prompts", "Configure Evals"];
  };

  const canProceedFromStep = (step: number): boolean => {
    switch (step) {
      case 0:
        // Experiment info step - need basic info, prompt versions, dataset, and evaluators
        // Only count already-added evaluators (not the current blank row)
        return !!(
          formData.name.trim() &&
          formData.promptVersions.length > 0 &&
          formData.datasetId &&
          formData.datasetVersion &&
          formData.evaluators.length > 0
        );
      case 1:
        // Configure prompt variables - need all mappings
        if (!formData.promptVariableMappings) return false;
        return promptVariables.every((varName) => !!formData.promptVariableMappings![varName]);
      case 2:
        // Configure evals - need all mappings for ALL evaluators
        if (!formData.evalVariableMappings || formData.evaluators.length === 0) return false;

        return formData.evaluators.every((evaluator) => {
          const evalMappings = formData.evalVariableMappings?.find((m) => m.evalName === evaluator.name && m.evalVersion === evaluator.version);
          if (!evalMappings) return false;

          const evalKey = `${evaluator.name}-${evaluator.version}`;
          const evalVars = evalVariables[evalKey]?.variables || [];
          return evalVars.every((varName) => !!evalMappings.mappings[varName]);
        });
      default:
        return false;
    }
  };

  const handleNext = async () => {
    if (currentStep === 0) {
      // Transition to next step immediately to prevent UI flicker
      setCompletedSteps((prev) => new Set(prev).add(currentStep));
      setCurrentStep((prev) => prev + 1);

      // Load prompt variables after moving to step 1
      // Always reload to reflect any changes made by going back and modifying step 0
      if (!taskId || !api || formData.promptVersions.length === 0) return;

      try {
        setLoadingPromptDetails(true);

        // Load variables from ALL selected prompt versions
        const allVariablesSet = new Set<string>();
        for (const promptVersion of formData.promptVersions) {
          const response = await api.api.getAgenticPromptApiV1TasksTaskIdPromptsPromptNameVersionsPromptVersionGet(
            promptVersion.promptName,
            String(promptVersion.version),
            taskId
          );
          const variables = response.data.variables || [];
          variables.forEach((v) => allVariablesSet.add(v));
        }

        const vars = Array.from(allVariablesSet);
        setPromptVariables(vars);

        // Initialize prompt variable mappings with auto-matching
        // Preserve existing mappings if they're still valid, otherwise auto-match
        if (datasetColumns.length > 0 && vars.length > 0) {
          const mappings: PromptVariableMappings = {};
          const existingMappings = formData.promptVariableMappings || {};

          vars.forEach((varName) => {
            // First check if we have an existing mapping that's still valid
            if (existingMappings[varName] && datasetColumns.includes(existingMappings[varName])) {
              mappings[varName] = existingMappings[varName];
            } else {
              // Otherwise try to auto-match
              const matchingColumn = datasetColumns.find((col) => col === varName);
              if (matchingColumn) {
                mappings[varName] = matchingColumn;
              }
            }
          });
          setFormData((prev) => ({ ...prev, promptVariableMappings: mappings }));
        } else {
          // Clear mappings if no variables
          setFormData((prev) => ({ ...prev, promptVariableMappings: {} }));
        }
      } catch (error) {
        console.error("Failed to load prompt variables:", error);
      } finally {
        setLoadingPromptDetails(false);
      }
    } else if (currentStep === 1) {
      // Load eval variables for first evaluator before moving to step 2
      // Always reload to reflect any changes made by going back and modifying evaluators
      if (formData.evaluators.length > 0) {
        const firstEval = formData.evaluators[0];
        await loadEvalVariablesForEvaluator(firstEval.name, firstEval.version);
        setCurrentEvalIndex(0);

        // Clean up eval variable mappings for removed evaluators
        // Only keep mappings for evaluators that still exist in formData.evaluators
        if (formData.evalVariableMappings) {
          const currentEvalKeys = new Set(formData.evaluators.map((e) => `${e.name}-${e.version}`));
          const cleanedMappings = formData.evalVariableMappings.filter((mapping) =>
            currentEvalKeys.has(`${mapping.evalName}-${mapping.evalVersion}`)
          );

          // Update formData with cleaned mappings
          if (cleanedMappings.length !== formData.evalVariableMappings.length) {
            setFormData((prev) => ({ ...prev, evalVariableMappings: cleanedMappings }));
          }
        }
      }
      setCompletedSteps((prev) => new Set(prev).add(currentStep));
      setCurrentStep((prev) => prev + 1);
    } else if (currentStep === 2) {
      // Within evals step, navigate between evaluators
      const nextEvalIndex = currentEvalIndex + 1;
      if (nextEvalIndex < formData.evaluators.length) {
        // Move to next evaluator within step 2
        const nextEval = formData.evaluators[nextEvalIndex];
        await loadEvalVariablesForEvaluator(nextEval.name, nextEval.version);
        setCurrentEvalIndex(nextEvalIndex);
      } else {
        // All evaluators configured, can submit
        setCompletedSteps((prev) => new Set(prev).add(currentStep));
      }
    }
  };

  const handleBack = () => {
    if (currentStep === 2 && currentEvalIndex > 0) {
      // Within evals step, go back to previous evaluator
      setCurrentEvalIndex((prev) => prev - 1);
    } else {
      // Go back to previous main step
      setCurrentStep((prev) => prev - 1);
      if (currentStep === 2) {
        // Reset eval index when leaving step 2
        setCurrentEvalIndex(0);
      }
    }
  };

  const isLastStep = () => {
    // We're at the last step when on step 2 and on the last evaluator
    return currentStep === 2 && currentEvalIndex === formData.evaluators.length - 1;
  };

  // Check if we can proceed from current eval in step 2
  const canProceedFromCurrentEval = (): boolean => {
    if (currentStep !== 2) return true;

    const evaluator = formData.evaluators[currentEvalIndex];
    if (!evaluator) return false;

    const evalMappings = formData.evalVariableMappings?.find((m) => m.evalName === evaluator.name && m.evalVersion === evaluator.version);
    if (!evalMappings) return false;

    const evalKey = `${evaluator.name}-${evaluator.version}`;
    const evalVars = evalVariables[evalKey]?.variables || [];
    return evalVars.every((varName) => !!evalMappings.mappings[varName]);
  };

  // Render step content
  const renderStepContent = () => {
    if (currentStep === 0) {
      return renderExperimentInfoStep();
    } else if (currentStep === 1) {
      return renderPromptVariableMappingStep();
    } else {
      return renderEvalVariableMappingStep(currentEvalIndex);
    }
  };

  const renderExperimentInfoStep = () => (
    <Box className="flex flex-col gap-4 mt-2">
      {/* Basic Info */}
      <TextField
        label="Experiment Name"
        value={formData.name}
        onChange={(e) => {
          setFormData((prev) => ({ ...prev, name: e.target.value }));
          if (errors.name) setErrors((prev) => ({ ...prev, name: undefined }));
        }}
        error={!!errors.name}
        helperText={errors.name}
        fullWidth
        required
        placeholder="e.g., Customer Support Tone Variations"
        autoFocus
      />

      <TextField
        label="Description"
        value={formData.description}
        onChange={(e) => setFormData((prev) => ({ ...prev, description: e.target.value }))}
        fullWidth
        multiline
        rows={2}
        placeholder="Describe the purpose of this experiment"
      />

      {/* Prompt Selection */}
      <Box sx={{ border: 1, borderColor: "divider", borderRadius: 1, p: 2 }}>
        <Box className="flex items-center gap-2 mb-2">
          <Typography variant="subtitle1" className="font-semibold">
            Prompt Versions *
          </Typography>
          <Tooltip
            title="Select one or more versions of a prompt to test. Each version will be evaluated against all test cases in your dataset."
            arrow
            placement="right"
          >
            <InfoOutlinedIcon
              sx={{
                fontSize: 18,
                color: "text.secondary",
                cursor: "help",
              }}
            />
          </Tooltip>
        </Box>

        <Box className="flex gap-2 mb-3 mt-4">
          <Autocomplete
            options={prompts}
            getOptionLabel={(option) => option.name}
            value={prompts.find((p) => p.name === selectedPromptName) || null}
            onChange={(_, value) => {
              setSelectedPromptName(value?.name || "");
              setPromptVersions([]);
              setVisibleOlderVersions([]);
            }}
            loading={loadingPrompts}
            renderInput={(params) => <TextField {...params} label="Select Prompt" placeholder="Search prompts..." />}
            className="flex-1"
          />
        </Box>

        {selectedPromptName && (
          <Box className="mb-3">
            <Typography variant="body2" className="text-gray-600 dark:text-gray-400 mb-2">
              Select versions to include (click to toggle):
            </Typography>
            {loadingPromptVersions ? (
              <CircularProgress size={24} />
            ) : (
              <>
                <Box className="flex flex-wrap gap-2 mb-3">
                  {(() => {
                    // Get the 5 most recent versions
                    const recentVersions = promptVersions.slice(0, 5);
                    // Get older versions that have been explicitly added
                    const olderVersionsToShow = promptVersions.filter(
                      (v) => visibleOlderVersions.includes(v.version) && !recentVersions.some((rv) => rv.version === v.version)
                    );
                    const allVisibleVersions = [...recentVersions, ...olderVersionsToShow];

                    return allVisibleVersions.map((version) => (
                      <Chip
                        key={version.version}
                        label={`v${version.version}`}
                        onClick={() => handleAddPromptVersion(version.version)}
                        color={
                          formData.promptVersions.some((pv) => pv.promptName === selectedPromptName && pv.version === version.version)
                            ? "primary"
                            : "default"
                        }
                        variant={
                          formData.promptVersions.some((pv) => pv.promptName === selectedPromptName && pv.version === version.version)
                            ? "filled"
                            : "outlined"
                        }
                      />
                    ));
                  })()}
                </Box>

                {promptVersions.length > 5 && (
                  <FormControl size="small" className="w-64">
                    <InputLabel>Add older version</InputLabel>
                    <Select
                      value=""
                      onChange={(e) => {
                        const version = Number(e.target.value);
                        if (version && !visibleOlderVersions.includes(version)) {
                          setVisibleOlderVersions((prev) => [...prev, version]);
                          handleAddPromptVersion(version);
                        }
                      }}
                      label="Add older version"
                    >
                      {promptVersions.slice(5).map((version) => (
                        <MenuItem key={version.version} value={version.version} disabled={visibleOlderVersions.includes(version.version)}>
                          v{version.version}
                        </MenuItem>
                      ))}
                    </Select>
                  </FormControl>
                )}
              </>
            )}
          </Box>
        )}

        {errors.promptVersions && (
          <Typography variant="caption" className="text-red-600 mt-2">
            {errors.promptVersions}
          </Typography>
        )}
      </Box>

      {/* Dataset Selection */}
      <Box sx={{ border: 1, borderColor: "divider", borderRadius: 1, p: 2 }}>
        <Box className="flex items-center gap-2 mb-2">
          <Typography variant="subtitle1" className="font-semibold">
            Dataset *
          </Typography>
          <Tooltip
            title="Choose the dataset and version to use for testing. Each row in the dataset will be used as input for the experiment."
            arrow
            placement="right"
          >
            <InfoOutlinedIcon
              sx={{
                fontSize: 18,
                color: "text.secondary",
                cursor: "help",
              }}
            />
          </Tooltip>
        </Box>

        <Box className="flex gap-2 mb-3 mt-4">
          <Autocomplete
            options={datasets}
            getOptionLabel={(option) => option.name}
            value={datasets.find((d) => d.id === formData.datasetId) || null}
            onChange={(_, value) => {
              setFormData((prev) => ({
                ...prev,
                datasetId: value?.id || "",
                datasetName: value?.name || "",
              }));
              if (value?.id) {
                loadDatasetVersions(value.id);
              } else {
                setDatasetVersions([]);
              }
              if (errors.datasetId) setErrors((prev) => ({ ...prev, datasetId: undefined }));
            }}
            loading={loadingDatasets}
            renderInput={(params) => <TextField {...params} label="Select Dataset" error={!!errors.datasetId} placeholder="Search datasets..." />}
            renderOption={(props, option) => (
              <li {...props}>
                <Box>
                  <Typography variant="body2">{option.name}</Typography>
                  {option.description && (
                    <Typography variant="caption" className="text-gray-600 dark:text-gray-400">
                      {option.description}
                    </Typography>
                  )}
                </Box>
              </li>
            )}
            className="flex-1"
          />

          <FormControl className="w-40">
            <InputLabel>Version</InputLabel>
            <Select
              value={formData.datasetVersion}
              onChange={(e) => {
                const versionNumber = e.target.value as number;
                setFormData((prev) => ({
                  ...prev,
                  datasetVersion: versionNumber,
                }));
                // Update columns when version changes
                const selectedVersion = datasetVersions.find((v) => v.version_number === versionNumber);
                if (selectedVersion) {
                  setDatasetColumns(selectedVersion.column_names);
                }
              }}
              label="Version"
              disabled={!formData.datasetId || loadingDatasetVersions}
              displayEmpty={false}
            >
              {datasetVersions.map((version) => (
                <MenuItem key={version.version_number} value={version.version_number}>
                  v{version.version_number}
                </MenuItem>
              ))}
            </Select>
          </FormControl>
        </Box>

        {(errors.datasetId || errors.datasetVersion) && (
          <Typography variant="caption" className="text-red-600">
            {errors.datasetId || errors.datasetVersion}
          </Typography>
        )}

        {/* Dataset Row Filter (Optional) */}
        {formData.datasetId && formData.datasetVersion && datasetColumns.length > 0 && (
          <Box className="mt-3 p-3 bg-gray-50 dark:bg-gray-800 border border-gray-200 dark:border-gray-700 rounded">
            <Box className="flex items-center gap-2 mb-2">
              <Typography variant="body2" className="font-medium text-gray-700 dark:text-gray-300">
                Filter Dataset Rows (Optional)
              </Typography>
              <Tooltip
                title="Optionally filter which dataset rows are included in this experiment. Only rows matching ALL specified conditions will be used."
                arrow
                placement="right"
              >
                <InfoOutlinedIcon
                  sx={{
                    fontSize: 16,
                    color: "text.secondary",
                    cursor: "help",
                  }}
                />
              </Tooltip>
            </Box>

            {formData.datasetRowFilter && formData.datasetRowFilter.length > 0 && (
              <Box className="flex flex-col gap-2 mb-2">
                {formData.datasetRowFilter.map((filter, index) => (
                  <Box key={index} className="flex gap-2 items-center">
                    <FormControl size="small" className="flex-1">
                      <InputLabel>Column</InputLabel>
                      <Select
                        value={filter.column_name}
                        onChange={(e) => {
                          const newFilters = [...(formData.datasetRowFilter || [])];
                          newFilters[index] = { ...filter, column_name: e.target.value };
                          setFormData((prev) => ({ ...prev, datasetRowFilter: newFilters }));
                        }}
                        label="Column"
                      >
                        {datasetColumns.map((column) => (
                          <MenuItem key={column} value={column}>
                            {column}
                          </MenuItem>
                        ))}
                      </Select>
                    </FormControl>
                    <TextField
                      size="small"
                      label="Value"
                      value={filter.column_value}
                      onChange={(e) => {
                        const newFilters = [...(formData.datasetRowFilter || [])];
                        newFilters[index] = { ...filter, column_value: e.target.value };
                        setFormData((prev) => ({ ...prev, datasetRowFilter: newFilters }));
                      }}
                      className="flex-1"
                    />
                    <Button
                      size="small"
                      onClick={() => {
                        const newFilters = formData.datasetRowFilter?.filter((_, i) => i !== index) || [];
                        setFormData((prev) => ({ ...prev, datasetRowFilter: newFilters }));
                      }}
                      color="error"
                    >
                      <CloseIcon fontSize="small" />
                    </Button>
                  </Box>
                ))}
              </Box>
            )}

            <Button
              size="small"
              variant="outlined"
              startIcon={<AddIcon />}
              onClick={() => {
                const newFilters = [...(formData.datasetRowFilter || []), { column_name: "", column_value: "" }];
                setFormData((prev) => ({ ...prev, datasetRowFilter: newFilters }));
              }}
            >
              Add Filter Condition
            </Button>
          </Box>
        )}
      </Box>

      {/* Evaluator Selection */}
      <Box sx={{ border: 1, borderColor: "divider", borderRadius: 1, p: 2 }}>
        <Box className="flex items-center gap-2 mb-2">
          <Typography variant="subtitle1" className="font-semibold">
            Evaluators *
          </Typography>
          <Tooltip
            title="Select one or more evaluators to assess the quality of the prompt outputs. Each evaluator will score the results based on its criteria."
            arrow
            placement="right"
          >
            <InfoOutlinedIcon
              sx={{
                fontSize: 18,
                color: "text.secondary",
                cursor: "help",
              }}
            />
          </Tooltip>
        </Box>

        <Box className="flex flex-col gap-3 mt-4">
          {/* Display selected evaluators as tiles */}
          {formData.evaluators.length > 0 && (
            <Box className="flex flex-wrap gap-2">
              {formData.evaluators.map((evaluator, index) => (
                <Chip
                  key={`${evaluator.name}-${evaluator.version}`}
                  label={`${evaluator.name} (v${evaluator.version})`}
                  onDelete={() => handleRemoveEvaluator(index)}
                  deleteIcon={<CloseIcon />}
                  color="primary"
                  variant="filled"
                />
              ))}
            </Box>
          )}

          {/* Evaluator selection form */}
          <Box className="flex gap-2 items-start">
            <FormControl className="flex-1">
              <InputLabel>Evaluator</InputLabel>
              <Select
                value={currentEvaluatorName}
                onChange={async (e) => {
                  const evalName = e.target.value;
                  setCurrentEvaluatorName(evalName);

                  if (evalName && !evaluatorVersions[evalName]) {
                    // Load versions - loadEvaluatorVersions will auto-set the most recent version
                    await loadEvaluatorVersions(evalName);
                  } else if (evalName && evaluatorVersions[evalName]) {
                    // If versions are already loaded, set the most recent (highest) version
                    const maxVersion = Math.max(...evaluatorVersions[evalName].map((v) => v.version));
                    setCurrentEvaluatorVersion(maxVersion);
                  } else {
                    setCurrentEvaluatorVersion("");
                  }
                }}
                label="Evaluator"
              >
                {evaluators.map((evaluator) => (
                  <MenuItem key={evaluator.name} value={evaluator.name}>
                    {evaluator.name}
                  </MenuItem>
                ))}
              </Select>
            </FormControl>

            <Box className="flex-1">
              <FormControl
                fullWidth
                error={
                  !!currentEvaluatorName &&
                  !!currentEvaluatorVersion &&
                  formData.evaluators.some((e) => e.name === currentEvaluatorName && e.version === currentEvaluatorVersion)
                }
              >
                <InputLabel>Version</InputLabel>
                <Select
                  value={currentEvaluatorVersion}
                  onChange={(e) => setCurrentEvaluatorVersion(e.target.value as number)}
                  label="Version"
                  disabled={!currentEvaluatorName}
                >
                  {currentEvaluatorName &&
                    evaluatorVersions[currentEvaluatorName]?.map((version) => (
                      <MenuItem key={version.version} value={version.version}>
                        v{version.version}
                      </MenuItem>
                    ))}
                </Select>
              </FormControl>
              {currentEvaluatorName &&
                currentEvaluatorVersion &&
                formData.evaluators.some((e) => e.name === currentEvaluatorName && e.version === currentEvaluatorVersion) && (
                  <Typography variant="caption" className="text-red-600 ml-3 mt-1">
                    This version has already been added
                  </Typography>
                )}
            </Box>

            <Button
              variant="outlined"
              onClick={handleAddEvaluator}
              disabled={
                !currentEvaluatorName ||
                !currentEvaluatorVersion ||
                formData.evaluators.some((e) => e.name === currentEvaluatorName && e.version === currentEvaluatorVersion)
              }
              startIcon={<AddIcon />}
              sx={{ height: 56, minWidth: 100 }}
            >
              Add
            </Button>
          </Box>
        </Box>

        {errors.evaluators && (
          <Typography variant="caption" className="text-red-600 mt-2">
            {errors.evaluators}
          </Typography>
        )}
      </Box>

      {errors.general && (
        <Typography variant="body2" className="text-red-600">
          {errors.general}
        </Typography>
      )}
    </Box>
  );

  const renderPromptVariableMappingStep = () => (
    <Box className="flex flex-col gap-4 mt-2">
      <Typography variant="body1" className="text-gray-700 dark:text-gray-300">
        Map each prompt variable to a dataset column. Variables that match column names exactly have been auto-filled.
      </Typography>

      {loadingPromptDetails ? (
        <Box className="flex justify-center p-4">
          <CircularProgress />
        </Box>
      ) : promptVariables.length === 0 ? (
        <Typography variant="body2" className="text-gray-600 dark:text-gray-400 italic">
          No variables found for this prompt.
        </Typography>
      ) : (
        <Box className="flex flex-col gap-3">
          {promptVariables.map((varName) => (
            <FormControl key={varName} fullWidth>
              <InputLabel required>{varName}</InputLabel>
              <Select
                value={formData.promptVariableMappings?.[varName] || ""}
                onChange={(e) => {
                  setFormData((prev) => ({
                    ...prev,
                    promptVariableMappings: {
                      ...prev.promptVariableMappings,
                      [varName]: e.target.value,
                    },
                  }));
                }}
                label={varName}
              >
                <MenuItem value="">
                  <em>Select a column</em>
                </MenuItem>
                {datasetColumns.map((column) => (
                  <MenuItem key={column} value={column}>
                    {column}
                  </MenuItem>
                ))}
              </Select>
            </FormControl>
          ))}
        </Box>
      )}
    </Box>
  );

  const renderEvalVariableMappingStep = (evalIndex: number) => {
    const evaluator = formData.evaluators[evalIndex];
    if (!evaluator) return null;

    const evalKey = `${evaluator.name}-${evaluator.version}`;
    const evalVars = evalVariables[evalKey]?.variables || [];
    const currentMappings = formData.evalVariableMappings?.find((m) => m.evalName === evaluator.name && m.evalVersion === evaluator.version);

    return (
      <Box className="flex flex-col gap-4 mt-2">
        <Box>
          <Typography variant="body2" className="text-gray-500 dark:text-gray-400 mb-1">
            Evaluator {evalIndex + 1} of {formData.evaluators.length}
          </Typography>
          <Typography variant="body1" className="text-gray-700 dark:text-gray-300 mb-2">
            Map each variable for{" "}
            <strong
              onClick={() => loadEvaluatorInstructions(evaluator.name, evaluator.version)}
              style={{
                cursor: "pointer",
                textDecoration: "underline",
                color: "#2563eb",
              }}
            >
              {evaluator.name} (v{evaluator.version})
            </strong>{" "}
            to either a dataset column or the experiment output.
          </Typography>
          <Box
            sx={(theme) => ({
              p: 1.5,
              bgcolor: alpha(theme.palette.info.main, 0.08),
              border: `1px solid ${alpha(theme.palette.info.main, 0.3)}`,
              borderRadius: 1,
            })}
          >
            <Typography variant="body2" sx={{ color: "text.primary" }}>
              <strong>Dataset Column:</strong> Use this when the evaluator needs information from your test data (e.g., expected answers, reference
              text, ground truth labels).
            </Typography>
            <Typography variant="body2" sx={{ color: "text.primary", mt: 0.5 }}>
              <strong>Experiment Output:</strong> Use this when the evaluator needs to assess the prompt's generated response (e.g., to check
              accuracy, relevance, or quality of the output).
            </Typography>
          </Box>
        </Box>

        {loadingEvalDetails ? (
          <Box className="flex justify-center p-4">
            <CircularProgress />
          </Box>
        ) : evalVars.length === 0 ? (
          <Typography variant="body2" className="text-gray-600 dark:text-gray-400 italic">
            No variables found for this evaluator.
          </Typography>
        ) : (
          <Box className="flex flex-col gap-4">
            {evalVars.map((varName) => {
              const mapping = currentMappings?.mappings[varName];
              const sourceType = mapping?.sourceType || "dataset_column";

              return (
                <Box key={varName} className="border border-gray-300 rounded p-3">
                  <Box className="flex items-center justify-between mb-3">
                    <Typography variant="subtitle2" className="font-medium">
                      {varName} *
                    </Typography>

                    <ToggleButtonGroup
                      value={sourceType}
                      exclusive
                      onChange={(event, newValue) => {
                        if (newValue === null) return;

                        const newMappings = formData.evalVariableMappings || [];
                        const existingIndex = newMappings.findIndex((m) => m.evalName === evaluator.name && m.evalVersion === evaluator.version);

                        const updatedMapping = {
                          evalName: evaluator.name,
                          evalVersion: evaluator.version,
                          mappings: {
                            ...(existingIndex >= 0 ? newMappings[existingIndex].mappings : {}),
                            [varName]:
                              newValue === "dataset_column"
                                ? {
                                    sourceType: "dataset_column" as VariableSourceType,
                                    datasetColumn: mapping?.datasetColumn || "",
                                  }
                                : {
                                    sourceType: "experiment_output" as VariableSourceType,
                                    jsonPath: mapping?.jsonPath || "",
                                  },
                          },
                        };

                        if (existingIndex >= 0) {
                          newMappings[existingIndex] = updatedMapping;
                        } else {
                          newMappings.push(updatedMapping);
                        }

                        setFormData((prev) => ({
                          ...prev,
                          evalVariableMappings: newMappings,
                        }));
                      }}
                      size="small"
                    >
                      <ToggleButton value="dataset_column">Dataset Column</ToggleButton>
                      <ToggleButton value="experiment_output">Experiment Output</ToggleButton>
                    </ToggleButtonGroup>
                  </Box>

                  {sourceType === "dataset_column" ? (
                    <Box>
                      <FormControl fullWidth size="small">
                        <InputLabel>Dataset Column</InputLabel>
                        <Select
                          value={mapping?.datasetColumn || ""}
                          onChange={(e) => {
                            const newMappings = formData.evalVariableMappings || [];
                            const existingIndex = newMappings.findIndex((m) => m.evalName === evaluator.name && m.evalVersion === evaluator.version);

                            const updatedMapping = {
                              evalName: evaluator.name,
                              evalVersion: evaluator.version,
                              mappings: {
                                ...(existingIndex >= 0 ? newMappings[existingIndex].mappings : {}),
                                [varName]: {
                                  sourceType: "dataset_column" as VariableSourceType,
                                  datasetColumn: e.target.value,
                                },
                              },
                            };

                            if (existingIndex >= 0) {
                              newMappings[existingIndex] = updatedMapping;
                            } else {
                              newMappings.push(updatedMapping);
                            }

                            setFormData((prev) => ({
                              ...prev,
                              evalVariableMappings: newMappings,
                            }));
                          }}
                          label="Dataset Column"
                        >
                          <MenuItem value="">
                            <em>Select a column</em>
                          </MenuItem>
                          {datasetColumns.map((column) => (
                            <MenuItem key={column} value={column}>
                              {column}
                            </MenuItem>
                          ))}
                        </Select>
                      </FormControl>
                    </Box>
                  ) : (
                    <Box className="p-3 bg-gray-50 dark:bg-gray-800 border border-gray-200 dark:border-gray-700 rounded">
                      <Typography variant="body2" className="text-gray-600 dark:text-gray-400 italic">
                        This variable will receive the full output from the prompt execution.
                      </Typography>
                    </Box>
                  )}
                </Box>
              );
            })}
          </Box>
        )}
      </Box>
    );
  };

  return (
    <>
      <Dialog open={open} onClose={handleCancel} maxWidth="md" fullWidth aria-labelledby="create-experiment-dialog-title">
        <DialogTitle id="create-experiment-dialog-title">{initialData ? "Create Experiment from Template" : "Create New Experiment"}</DialogTitle>
        <DialogContent>
          {isLoadingInitialData || isInitializing ? (
            <Box className="flex justify-center items-center py-8">
              <CircularProgress />
            </Box>
          ) : (
            <>
              <Box className="mb-4">
                <Stepper activeStep={currentStep} alternativeLabel>
                  {getStepLabels().map((label, index) => (
                    <Step key={label} completed={completedSteps.has(index)}>
                      <StepLabel>{label}</StepLabel>
                    </Step>
                  ))}
                </Stepper>
              </Box>

              {renderStepContent()}
            </>
          )}
        </DialogContent>
        <DialogActions className="px-6 pb-4">
          <Button onClick={handleCancel} disabled={isSubmitting}>
            Cancel
          </Button>
          {(currentStep > 0 || currentEvalIndex > 0) && (
            <Button onClick={handleBack} disabled={isSubmitting}>
              Back
            </Button>
          )}
          {!isLastStep() ? (
            <Button
              onClick={handleNext}
              variant="contained"
              color="primary"
              disabled={(currentStep === 2 ? !canProceedFromCurrentEval() : !canProceedFromStep(currentStep)) || isSubmitting}
            >
              {currentStep === 0 ? "Configure Prompts" : currentStep === 1 ? "Configure Evals" : "Next Evaluator"}
            </Button>
          ) : (
            <Button
              onClick={handleSubmit}
              variant="contained"
              color="primary"
              disabled={!canProceedFromCurrentEval() || isSubmitting}
              startIcon={isSubmitting ? <CircularProgress size={16} /> : null}
            >
              {isSubmitting ? "Creating..." : "Create Experiment"}
            </Button>
          )}
        </DialogActions>
      </Dialog>

      {/* Evaluator Instructions Modal */}
      <Dialog
        open={instructionsModalOpen}
        onClose={() => setInstructionsModalOpen(false)}
        maxWidth="lg"
        fullWidth
        aria-labelledby="evaluator-instructions-dialog-title"
      >
        <DialogTitle id="evaluator-instructions-dialog-title">
          {selectedEvalInstructions && (
            <>
              {selectedEvalInstructions.name} (v{selectedEvalInstructions.version}) - Instructions
            </>
          )}
        </DialogTitle>
        <DialogContent sx={{ pb: 1 }}>
          {loadingInstructions ? (
            <Box className="flex justify-center p-4">
              <CircularProgress />
            </Box>
          ) : selectedEvalInstructions ? (
            <Box
              className="whitespace-pre-wrap font-mono text-sm p-4 bg-gray-50 dark:bg-gray-800 border border-gray-200 dark:border-gray-700 rounded"
              sx={{
                maxHeight: "70vh",
                overflowY: "auto",
                overflowX: "auto",
              }}
            >
              {selectedEvalInstructions.instructions || "No instructions available for this evaluator."}
            </Box>
          ) : null}
        </DialogContent>
        <DialogActions>
          <Button onClick={() => setInstructionsModalOpen(false)} color="primary">
            Close
          </Button>
        </DialogActions>
      </Dialog>

      {/* Success/Error Toast */}
      <Snackbar {...snackbarProps}>
        <Alert {...alertProps} />
      </Snackbar>
    </>
  );
};
