import { useStatsigClient } from '@statsig/react-bindings';
import { isEqual, memoize } from 'lodash-es';
import { getIn, setIn } from 'lodash-redux-immutability';
import { usePathname } from 'next/navigation';
import {
  Dispatch,
  SetStateAction,
  useCallback,
  useEffect,
  useMemo,
  useRef,
  useState,
} from 'react';

import { useStores } from '@/app/(root)/AppProviders';
import applyPersona from '@/components/clipBrowser/actions/usePersona';
import { ModalTypes } from '@/components/modal/constants/ModalTypes';
import renderAuxiliaryClip from '@/components/studio/apiWrappers/renderAuxiliaryClip';
import DEFAULT_STATE from '@/components/studio/defaultState';
import { StudioProjectState } from '@/components/studio/types';
import { toast } from '@/components/toast/Toast';
import { useFetchClip } from '@/hooks/useClip';
import { createContext } from '@/hooks/useContextSelector';
import { useStateMemorization } from '@/hooks/useStateMemorization';
import { useApiClient } from '@/lib/apiClient';
import logWebUserEvent from '@/logging/logWebUserEvent';
import { Clip, Playlist } from '@/state/clipStore';
import { PersonaMetadata } from '@/state/personaStore';
import { FeatureKey, PlanFeature } from '@/state/sessionStore';
import { MIN_EXTEND_FROM_SECONDS } from '@/utils/constants';
import { isFeatureEnabledForPlan } from '@/utils/session';
import { getDefaultModel } from '@/utils/utils';

import { setOneUntimedCondition } from './actions/currentModeSetters';
import { PATHS_TO_RESET } from './actions/resetCreateInputs';
import defaultCreateFormState from './defaultCreateFormState';
import {
  AddMenuStates,
  ConditionTypes,
  CreateFormModals,
  CreateFormState,
  CreateModes,
  ImposedStateChanges,
  LyricsInputModes,
  PlaylistCondition,
} from './types';
import { AddConditionType, useConditionAdder } from './useConditionAdder';
import useCreateV2Generate from './useCreateV2Generate';
import useCredits from './useCredits';
import useModelSwitchOnUpgrade from './useModelSwitchOnUpgrade';
import useModelTypeOptions from './useModelTypeOptions';
import usePlaylistConditionAccess from './usePlaylistConditionAccess';
import { usePrompts } from './usePromptsV2';
import useValidateStateChange from './useValidateStateChange';

// Re-export for convenience
export { AddConditionType } from './useConditionAdder';

const confirmSideEffects = false;

// Define which paths should be automatically memorized
const pathsToMemorize: (string | number)[][] = [
  ['global', 'model'],
  ['global', 'mode'],
  ['global', 'lyricsModel'],
];

export const useCreateForm = () => {
  const isStudio = usePathname().startsWith('/studio');
  const modelTypeOptions = useModelTypeOptions();
  const { session, menus } = useStores();
  const statsigClient = useStatsigClient();

  // Create base default state
  const baseDefaultState = useMemo(() => {
    let result = defaultCreateFormState;
    result = setIn(
      result,
      ['global', 'model'],
      getDefaultModel(modelTypeOptions)
    );
    result = setIn(
      result,
      ['global', 'mode'],
      isStudio ? CreateModes.CUSTOM : CreateModes.SIMPLE
    );
    return result;
  }, [modelTypeOptions, isStudio]);

  // Use single instance of the memorization hook
  const {
    applyMemorizedValuesToDefaultState,
    onStateChange: memorizeOnStateChange,
    getMemorizedStateFromPath,
  } = useStateMemorization<CreateFormState>(
    isStudio
      ? 'studio-create-form-memorized-state-sep-1-2025'
      : 'create-form-memorized-state-sep-1-2025',
    pathsToMemorize
  );

  const hasMemorizedModel = !!getMemorizedStateFromPath(['global', 'model']);

  const hasInitialized = modelTypeOptions.length > 0 || hasMemorizedModel;
  const lastHasInitialized = useRef(hasInitialized);

  // Create initial state with memorized values
  const initialState = useMemo(() => {
    return applyMemorizedValuesToDefaultState(baseDefaultState);
  }, [applyMemorizedValuesToDefaultState, baseDefaultState, pathsToMemorize]);

  const [state, setState] = useState<CreateFormState>(initialState);
  const stateRef = useRef<CreateFormState>(state);

  const setStateAndRefIfUnlocked = useCallback(
    (newState: CreateFormState) => {
      if (!hasInitialized) return;
      stateRef.current = newState;
      setState(newState);
    },
    [hasInitialized]
  );

  useEffect(() => {
    if (hasInitialized && !lastHasInitialized.current) {
      setStateAndRefIfUnlocked(initialState);
    }
    lastHasInitialized.current = hasInitialized;
  }, [setStateAndRefIfUnlocked, hasInitialized, initialState]);

  // Automatically memorize changes to specified paths
  useEffect(() => {
    if (!hasInitialized) return;
    memorizeOnStateChange(state);
  }, [state, memorizeOnStateChange, hasInitialized]);

  const { generate, isGenerating, createBlockingErrorMessage } =
    useCreateV2Generate(state);

  const validateStateChange = useValidateStateChange();

  const [nextState, setNextStateChange] = useState<CreateFormState | null>(
    null
  );

  const [validationImposedStateChanges, setValidationImposedStateChanges] =
    useState<ImposedStateChanges[]>([]);

  const [
    isConfirmingValidationImposedStateChanges,
    setIsConfirmingValidationImposedStateChanges,
  ] = useState<boolean>(false);

  const [isAddMenuOpen, setIsAddMenuOpen] = useState<boolean>(false);
  const [addMenuState, setAddMenuState] = useState<AddMenuStates>(
    AddMenuStates.MAIN
  );
  const [activeModal, setActiveModal] = useState<CreateFormModals | null>(null);
  const [isCreatePersonaModalOpen, setIsCreatePersonaModalOpen] =
    useState<boolean>(false);
  const [isSettingsMenuOpen, setIsSettingsMenuOpen] = useState<boolean>(false);

  // Pending persona for overwrite modal (kept separate due to different structure)
  const [pendingPersona, setPendingPersona] = useState<{
    persona: PersonaMetadata & { clip?: { image_url?: string | null } };
    rootClip: Clip;
  } | null>(null);

  // Studio timeline rendering state
  const [isStudioTimelineRendering, setIsStudioTimelineRendering] =
    useState<boolean>(false);
  const [
    lastStudioTimelineRenderTimestamp,
    setLastStudioTimelineRenderTimestamp,
  ] = useState<number>(Date.now());

  const isRemix = useMemo(() => {
    return (
      state.global.mode === CreateModes.CUSTOM &&
      (!!state.global.untimedConditions[ConditionTypes.COVER] ||
        !!(
          state.global.timedCondition?.type === ConditionTypes.EXTEND &&
          state.global.timedCondition.condition?.clipId
        ) ||
        !!state[CreateModes.CUSTOM].reusedFromClipId)
    );
  }, [state]);

  const {
    lyricsPrompts,
    tagsPrompts,
    fetchNextLyricsPage,
    fetchNextTagsPage,
    isLoadingLyrics,
    isLoadingTags,
    saveLyricsPrompt,
    saveTagsPrompt,
    lastSavedLyricsPromptState,
    lastSavedTagsPromptState,
    deletePromptFromCache,
  } = usePrompts();

  const applyNextStateChange = useCallback(() => {
    if (nextState) setState(nextState);
    setNextStateChange(null);
    setValidationImposedStateChanges([]);
    setIsConfirmingValidationImposedStateChanges(false);
  }, [nextState]);

  const getPromptToSave = useCallback(() => {
    const promptToSave = {
      gptDescriptionPrompt:
        state[CreateModes.CUSTOM].lyricsMode === LyricsInputModes.AUTO &&
        state[CreateModes.CUSTOM].lyricsPrompt !== ''
          ? state[CreateModes.CUSTOM].lyricsPrompt
          : undefined,
      tags:
        state[CreateModes.CUSTOM].styles !== ''
          ? state[CreateModes.CUSTOM].styles
          : undefined,
      negativeTags:
        state[CreateModes.CUSTOM].excludeStyles !== ''
          ? state[CreateModes.CUSTOM].excludeStyles
          : undefined,
      lyrics:
        state[CreateModes.CUSTOM].lyrics !== ''
          ? state[CreateModes.CUSTOM].lyrics
          : undefined,
      title:
        state[CreateModes.CUSTOM].title !== ''
          ? state[CreateModes.CUSTOM].title
          : undefined,
    };
    return promptToSave;
  }, [state]);

  const validateAndSetState = useCallback(
    (updater: React.SetStateAction<CreateFormState>) => {
      const prevState = stateRef.current;
      const newState =
        typeof updater === 'function'
          ? (updater as (prev: CreateFormState) => CreateFormState)(prevState)
          : updater;

      const { nextState, imposedStateChanges, errors } = validateStateChange(
        prevState,
        newState
      );

      if (errors.length) {
        errors.forEach((error) => {
          toast({
            title: error,
            status: 'error',
          });
        });
      }

      if (
        imposedStateChanges.length &&
        prevState.global.mode !== CreateModes.CHAT
      ) {
        imposedStateChanges.forEach((change) => {
          switch (change) {
            case ImposedStateChanges.EXTEND_TIME_SET_TO_MIN:
              toast({
                title: '"Extend from" updated',
                description: `Extend-from time was set to the minimum value of ${MIN_EXTEND_FROM_SECONDS} second${MIN_EXTEND_FROM_SECONDS === 1 ? '' : 's'}.`,
                status: 'info',
                duration: 3000,
                isClosable: true,
              });
              break;
            case ImposedStateChanges.REMOVE_PLAYLIST_CONDITION:
              toast({
                title: 'Inspiration removed',
                description:
                  'Inspiration was removed due to a conflict with another condition.',
                status: 'info',
                duration: 3000,
                isClosable: true,
              });
              break;
            case ImposedStateChanges.REMOVE_TIMED_CONDITION:
              toast({
                title: 'Vocals/Instrumentals removed',
                description:
                  'Vocals or instrumentals were removed due to a conflict with inspiration.',
                status: 'info',
                duration: 3000,
                isClosable: true,
              });
              break;
            case ImposedStateChanges.REMOVE_COVER_CONDITION:
              toast({
                title: 'Cover condition removed',
                description:
                  'Cover condition was removed due to a conflict with another condition.',
                status: 'info',
                duration: 3000,
                isClosable: true,
              });
              break;
            case ImposedStateChanges.REMOVE_PERSONA:
              toast({
                title: 'Persona removed',
                description:
                  'Persona was removed due to a conflict with another condition.',
                status: 'info',
                duration: 3000,
                isClosable: true,
              });
              break;
            case ImposedStateChanges.MODEL_CHANGE:
              toast({
                title: 'Model changed',
                description:
                  'Model was automatically changed to support your selected conditions.',
                status: 'info',
                duration: 3000,
                isClosable: true,
              });
              break;
            case ImposedStateChanges.TRUNCATE_INPUTS:
              toast({
                title: 'Inputs truncated',
                description:
                  'Some inputs were automatically shortened to fit the selected model.',
                status: 'info',
                duration: 3000,
                isClosable: true,
              });
              break;
          }
        });
      }

      if (imposedStateChanges.length && confirmSideEffects) {
        setNextStateChange(nextState);
        setValidationImposedStateChanges(imposedStateChanges);
        setIsConfirmingValidationImposedStateChanges(true);
      } else {
        stateRef.current = nextState;
        setStateAndRefIfUnlocked(nextState);
      }
    },

    [setStateAndRefIfUnlocked, validateStateChange]
  );

  // Use the condition adder hook
  const {
    addCondition,
    pendingCondition,
    handleConditionOverwriteConfirm,
    handleConditionOverwriteCancel,
    handleConditionOverwriteClose,
  } = useConditionAdder(validateAndSetState, stateRef, setActiveModal);

  const { creditsRemaining, chargeCredits } = useCredits();
  const { planAllowsPlaylistCondition, showPlaylistConditionUpsellOrRedirect } =
    usePlaylistConditionAccess();

  const fetchClip = useFetchClip();

  const checkCanAccessPersonas = useCallback(() => {
    if (isFeatureEnabledForPlan(session, PlanFeature.Persona)) {
      return true;
    }

    menus.setCurrentUpsellFeature(FeatureKey.PERSONAS);
    menus.openModal(ModalTypes.UPSELL_MODAL);

    return false;
  }, [session, statsigClient, menus]);

  const checkCanAccessPlaylists = useCallback(() => {
    if (planAllowsPlaylistCondition) return true;
    showPlaylistConditionUpsellOrRedirect();
    return false;
  }, [planAllowsPlaylistCondition, showPlaylistConditionUpsellOrRedirect]);

  const showPersonaConflictToasts = useCallback(
    (
      hasPlaylist: boolean,
      hasVocalsOrInstrumentals: boolean,
      conditionType?: ConditionTypes.OVERPAINT | ConditionTypes.UNDERPAINT
    ) => {
      const canUseVoxPersona = session.isFeatureAllowed('vox-persona').enabled;

      if (hasPlaylist) {
        toast({
          title: 'Inspiration removed',
          description:
            'Inspiration was removed to add Persona.' +
            (canUseVoxPersona
              ? ' To use Vox Persona Inspiration, add the Persona and enable Vox Persona toggle before adding Inspiration.'
              : ''),
          status: 'info',
          duration: canUseVoxPersona ? 10000 : 3000,
          isClosable: true,
        });
      }

      if (hasVocalsOrInstrumentals && conditionType) {
        const conditionName =
          conditionType === ConditionTypes.OVERPAINT
            ? 'Add Vocals'
            : 'Add Instrumental';
        toast({
          title: 'Condition removed',
          description: `${conditionName} was removed to add Persona.`,
          status: 'info',
          duration: 3000,
          isClosable: true,
        });
      }
    },
    [session]
  );

  const attemptOpenPersonaModal = useCallback(() => {
    if (!checkCanAccessPersonas()) return;
    setActiveModal(CreateFormModals.Personas);
  }, [checkCanAccessPersonas]);

  const attemptOpenInspoModal = useCallback(() => {
    if (!checkCanAccessPlaylists()) return;
    setActiveModal(CreateFormModals.Inspo);
  }, [checkCanAccessPlaylists]);

  const addPersona = useCallback(
    async (
      persona: PersonaMetadata & { clip?: { image_url?: string | null } },
      overwriteStyles?: boolean
    ) => {
      if (!persona.root_clip_id) return;

      const rootClip = await fetchClip(persona.root_clip_id);
      if (!rootClip) return;

      const currentState = stateRef.current;
      // Check if styles are not empty and we need to ask for overwrite
      const currentStyles = currentState[CreateModes.CUSTOM].styles || '';
      if (currentStyles.trim() !== '' && overwriteStyles === undefined) {
        setPendingPersona({ persona, rootClip });
        setActiveModal(CreateFormModals.OverwriteStyles);
        logWebUserEvent({
          actionName: 'PersonaOverwriteModalShown',
          context: {
            personaId: persona.id || undefined,
            rootClipId: persona.root_clip_id || rootClip.id || undefined,
          },
        });
        return;
      }

      // Apply persona directly (no modal needed) - show toasts and apply
      const hasPlaylist = !!currentState.global.untimedConditions.playlist;
      const hasVocalsOrInstrumentals: boolean = !!(
        currentState.global.timedCondition &&
        (currentState.global.timedCondition.type === ConditionTypes.OVERPAINT ||
          currentState.global.timedCondition.type === ConditionTypes.UNDERPAINT)
      );
      const conditionType = currentState.global.timedCondition?.type as
        | ConditionTypes.OVERPAINT
        | ConditionTypes.UNDERPAINT
        | undefined;

      showPersonaConflictToasts(
        hasPlaylist,
        hasVocalsOrInstrumentals,
        conditionType
      );

      // Apply persona with or without overwrite
      validateAndSetState((prev: CreateFormState) => {
        const stateWithoutPlaylistAndTimed = {
          ...prev,
          global: {
            ...prev.global,
            timedCondition: hasVocalsOrInstrumentals
              ? null
              : prev.global.timedCondition,
            untimedConditions: {
              ...prev.global.untimedConditions,
              playlist: undefined,
            },
          },
        };

        return applyPersona(persona, rootClip, {
          overwriteStyles: overwriteStyles === true,
        })(stateWithoutPlaylistAndTimed);
      });
    },
    [validateAndSetState, fetchClip, session, showPersonaConflictToasts]
  );

  const applyPendingPersona = useCallback(
    (overwriteStyles: boolean) => {
      if (!pendingPersona) return;
      const { persona, rootClip } = pendingPersona;
      const currentState = stateRef.current;
      const hasPlaylist = !!currentState.global.untimedConditions.playlist;
      const hasVocalsOrInstrumentals: boolean = !!(
        currentState.global.timedCondition &&
        (currentState.global.timedCondition.type === ConditionTypes.OVERPAINT ||
          currentState.global.timedCondition.type === ConditionTypes.UNDERPAINT)
      );
      const conditionType = currentState.global.timedCondition?.type as
        | ConditionTypes.OVERPAINT
        | ConditionTypes.UNDERPAINT
        | undefined;

      // Show toasts when persona is actually being applied (after modal confirmation)
      showPersonaConflictToasts(
        hasPlaylist,
        hasVocalsOrInstrumentals,
        conditionType
      );

      setPendingPersona(null);
      setActiveModal(null);

      validateAndSetState((prev: CreateFormState) => {
        const stateWithoutPlaylistAndTimed = {
          ...prev,
          global: {
            ...prev.global,
            timedCondition: hasVocalsOrInstrumentals
              ? null
              : prev.global.timedCondition,
            untimedConditions: {
              ...prev.global.untimedConditions,
              playlist: undefined,
            },
          },
        };

        return applyPersona(persona, rootClip, {
          overwriteStyles,
        })(stateWithoutPlaylistAndTimed);
      });
    },
    [pendingPersona, validateAndSetState, session, showPersonaConflictToasts]
  );

  const handleOverwriteConfirm = useCallback(() => {
    if (!pendingPersona) return;
    logWebUserEvent({
      actionName: 'PersonaOverwriteConfirmed',
      context: {
        personaId: pendingPersona.persona.id || undefined,
        rootClipId: pendingPersona.rootClip.id || undefined,
      },
    });
    applyPendingPersona(true);
  }, [applyPendingPersona, pendingPersona]);

  const handleOverwriteCancel = useCallback(() => {
    if (!pendingPersona) return;
    logWebUserEvent({
      actionName: 'PersonaOverwriteCancelled',
      context: {
        personaId: pendingPersona.persona.id || undefined,
        rootClipId: pendingPersona.rootClip.id || undefined,
      },
    });
    applyPendingPersona(false);
  }, [applyPendingPersona, pendingPersona]);

  const handleOverwriteClose = useCallback(() => {
    setPendingPersona(null);
    setActiveModal(null);
  }, [setActiveModal]);

  const addInspoPlaylist = useCallback(
    async (playlist: Playlist) => {
      if (!playlist.id) return;

      const currentState = stateRef.current;
      const hasPersona =
        !!currentState.global.untimedConditions.persona?.length;
      const hasExistingInspiration =
        !!currentState.global.untimedConditions[ConditionTypes.INSPIRATION]
          ?.clipIds?.length;
      const hasExistingTitle = !!currentState[CreateModes.CUSTOM].title?.length;
      const isVoxPersonaEnabled =
        !!currentState.global.untimedConditions[ConditionTypes.VOX_PERSONA]
          ?.enabled;

      const canUseVoxPersona = session.isFeatureAllowed('vox-persona').enabled;

      // Only remove persona if vox persona is not enabled (to allow VOX_PLAYLIST_CONDITION)
      if (hasPersona && !isVoxPersonaEnabled) {
        toast({
          title: 'Persona removed',
          description:
            'Persona was removed to add Inspiration.' +
            (canUseVoxPersona
              ? ' To use Vox Persona Inspiration, enable Vox Persona toggle before adding Inspiration.'
              : ''),
          status: 'info',
          duration: canUseVoxPersona ? 10000 : 3000,
          isClosable: true,
        });
      }

      if (hasExistingInspiration) {
        toast({
          title: 'Inspiration clips cleared',
          description:
            'Individual inspiration clips were cleared to add playlist.',
          status: 'info',
          duration: 3000,
          isClosable: true,
        });
      }

      validateAndSetState((prev: CreateFormState) => {
        const existingTimedCondition = prev.global.timedCondition;

        if (existingTimedCondition) {
          let description = '';
          switch (existingTimedCondition.type) {
            case ConditionTypes.EXTEND:
              description = 'Extend condition was removed to add inspiration.';
              break;
            case ConditionTypes.OVERPAINT:
              description = 'Add Vocals was removed to add inspiration.';
              break;
            case ConditionTypes.UNDERPAINT:
              description = 'Add Instrumental was removed to add inspiration.';
              break;
            default:
              description =
                'Previous condition was removed to add inspiration.';
          }

          toast({
            title: 'Condition removed',
            description,
            status: 'info',
            duration: 3000,
            isClosable: true,
          });
        }

        const stateWithoutTimedCondition = {
          ...prev,
          global: {
            ...prev.global,
            timedCondition: null,
          },
        };

        const stateWithoutPersonasAndInspiration = {
          ...stateWithoutTimedCondition,
          global: {
            ...stateWithoutTimedCondition.global,
            untimedConditions: {
              ...stateWithoutTimedCondition.global.untimedConditions,
              // Only remove persona if vox persona is not enabled (to allow VOX_PLAYLIST_CONDITION)
              persona: isVoxPersonaEnabled
                ? stateWithoutTimedCondition.global.untimedConditions.persona
                : undefined,
              [ConditionTypes.INSPIRATION]: undefined,
            },
          },
        };

        // Only set title if user hasn't already set a custom title
        const stateWithPlaylist = hasExistingTitle
          ? stateWithoutPersonasAndInspiration
          : {
              ...stateWithoutPersonasAndInspiration,
              [CreateModes.CUSTOM]: {
                ...stateWithoutPersonasAndInspiration[CreateModes.CUSTOM],
                title: `Inspired by ${playlist.name || 'playlist'}`,
              },
            };

        return setOneUntimedCondition(ConditionTypes.PLAYLIST, {
          playlistId: playlist.id,
        } as PlaylistCondition)(stateWithPlaylist);
      });
    },
    [validateAndSetState]
  );

  const addInspirationClips = useCallback(
    (clipIds: string[]) => {
      if (!clipIds.length) return;

      validateAndSetState((prevState) => {
        // Get current inspiration clip IDs
        const currentInspirationClips =
          prevState.global.untimedConditions[ConditionTypes.INSPIRATION]
            ?.clipIds || [];

        // Add new clips, filter duplicates, and limit to 4 total
        const allClipIds = [...currentInspirationClips, ...clipIds];
        const uniqueClipIds = Array.from(new Set(allClipIds)).slice(0, 4);

        const isVoxPersonaEnabled =
          !!prevState.global.untimedConditions[ConditionTypes.VOX_PERSONA]
            ?.enabled;

        // Clear any existing timed conditions and personas when adding inspiration
        // Only remove persona if vox persona is not enabled (to allow VOX_PLAYLIST_CONDITION)
        const stateWithoutTimedCondition = {
          ...prevState,
          global: {
            ...prevState.global,
            timedCondition: null,
            untimedConditions: {
              ...prevState.global.untimedConditions,
              persona: isVoxPersonaEnabled
                ? prevState.global.untimedConditions.persona
                : undefined,
              [ConditionTypes.INSPIRATION]: { clipIds: uniqueClipIds },
            },
          },
        };

        return stateWithoutTimedCondition;
      });
    },
    [validateAndSetState]
  );

  const removeInspirationClip = useCallback(
    (clipIdToRemove: string) => {
      validateAndSetState((prevState) => {
        const currentInspirationClips =
          prevState.global.untimedConditions[ConditionTypes.INSPIRATION]
            ?.clipIds || [];
        const filteredClips = currentInspirationClips.filter(
          (id) => id !== clipIdToRemove
        );

        return {
          ...prevState,
          global: {
            ...prevState.global,
            untimedConditions: {
              ...prevState.global.untimedConditions,
              [ConditionTypes.INSPIRATION]:
                filteredClips.length > 0
                  ? { clipIds: filteredClips }
                  : undefined,
            },
          },
        };
      });
    },
    [validateAndSetState]
  );

  const clearInspirationClips = useCallback(() => {
    validateAndSetState((prevState) => ({
      ...prevState,
      global: {
        ...prevState.global,
        untimedConditions: {
          ...prevState.global.untimedConditions,
          [ConditionTypes.INSPIRATION]: undefined,
        },
      },
    }));
  }, [validateAndSetState]);

  const addSampleClips = useCallback(
    (clipIds: string[]) => {
      if (!clipIds.length) return;

      validateAndSetState((prevState) => {
        // Get current sample clip IDs
        const currentSampleClips =
          prevState.global.untimedConditions[ConditionTypes.SAMPLE]?.clipIds ||
          [];

        // Add new clips, filter duplicates, and limit to 10 total
        const allClipIds = [...currentSampleClips, ...clipIds];
        const uniqueClipIds = Array.from(new Set(allClipIds)).slice(0, 10);

        // Clear any existing timed conditions and personas when adding sample
        const stateWithoutTimedCondition = {
          ...prevState,
          global: {
            ...prevState.global,
            timedCondition: null,
            untimedConditions: {
              ...prevState.global.untimedConditions,
              [ConditionTypes.SAMPLE]: { clipIds: uniqueClipIds },
            },
          },
        };

        return stateWithoutTimedCondition;
      });
    },
    [validateAndSetState]
  );

  const removeSampleClip = useCallback(
    (clipIdToRemove: string) => {
      validateAndSetState((prevState) => {
        const currentSampleClips =
          prevState.global.untimedConditions[ConditionTypes.SAMPLE]?.clipIds ||
          [];
        const filteredClips = currentSampleClips.filter(
          (id) => id !== clipIdToRemove
        );

        return {
          ...prevState,
          global: {
            ...prevState.global,
            untimedConditions: {
              ...prevState.global.untimedConditions,
              [ConditionTypes.SAMPLE]:
                filteredClips.length > 0
                  ? { clipIds: filteredClips }
                  : undefined,
            },
          },
        };
      });
    },
    [validateAndSetState]
  );

  const clearSampleClips = useCallback(() => {
    validateAndSetState((prevState) => ({
      ...prevState,
      global: {
        ...prevState.global,
        untimedConditions: {
          ...prevState.global.untimedConditions,
          [ConditionTypes.SAMPLE]: undefined,
        },
      },
    }));
  }, [validateAndSetState]);

  const apiClient = useApiClient();

  const addStudioProjectStateCondition = useCallback(
    async (
      state: StudioProjectState,
      conditionType: ConditionTypes,
      options: {
        projectId?: string;
        studioProjectId: string;
        startBeats: number;
        endBeats: number;
        title?: string;
        lyrics?: string;
      }
    ) => {
      try {
        if (
          ![
            ConditionTypes.COVER,
            ConditionTypes.EXTEND,
            ConditionTypes.UNDERPAINT,
            ConditionTypes.OVERPAINT,
          ].includes(conditionType)
        ) {
          throw new Error(`Invalid condition type: ${conditionType}`);
        }

        setIsStudioTimelineRendering(true);
        setLastStudioTimelineRenderTimestamp(Date.now());
        validateAndSetState((prev) =>
          setIn(prev, ['global', 'mode'], CreateModes.CUSTOM)
        );

        const clip = await renderAuxiliaryClip({
          title: options.title || 'Studio Selection',
          state: {
            ...DEFAULT_STATE,
            tracks: state.tracks,
            timing: state.timing,
          },
          apiClient,
          projectId: options.projectId,
          fromStudioProjectId: options.studioProjectId,
          startBeats: options.startBeats,
          endBeats: options.endBeats,
          lyrics: options.lyrics || '',
        });

        await fetchClip(clip.id);

        // Map ConditionTypes to AddConditionType
        let addConditionType: AddConditionType;
        switch (conditionType) {
          case ConditionTypes.COVER:
            addConditionType = AddConditionType.COVER;
            break;
          case ConditionTypes.EXTEND:
            addConditionType = AddConditionType.EXTEND;
            break;
          case ConditionTypes.UNDERPAINT:
            addConditionType = AddConditionType.UNDERPAINT;
            break;
          case ConditionTypes.OVERPAINT:
            addConditionType = AddConditionType.OVERPAINT;
            break;
          default:
            throw new Error(`Invalid condition type: ${conditionType}`);
        }

        await addCondition(clip.id, addConditionType);
      } catch (error) {
        console.error(error);
        toast({
          title: 'Error preparing selection for edit',
          description: 'Please try again.',
          status: 'error',
        });
        throw error; // Re-throw the error so the finally block can still execute
      } finally {
        setIsStudioTimelineRendering(false);
      }
    },
    [apiClient, validateAndSetState, fetchClip, addCondition]
  );

  const setChameleonMode = useCallback(
    (enabled: boolean) => {
      validateAndSetState((prevState) => {
        if (enabled) {
          return setOneUntimedCondition(ConditionTypes.CHAMELEON, {
            enabled: true,
          })(prevState);
        } else {
          const { chameleon, ...otherConditions } =
            prevState.global.untimedConditions;
          return setIn(
            prevState,
            ['global', 'untimedConditions'],
            otherConditions
          );
        }
      });
    },
    [validateAndSetState]
  );

  const getChameleonMode = useCallback(() => {
    const chameleonCondition = state.global.untimedConditions.chameleon;
    return chameleonCondition?.enabled ?? false;
  }, [state]);

  const setChameleonTargetUser = useCallback(
    (
      user: {
        userId: string;
        handle: string;
        displayName?: string | null;
        avatarImageUrl?: string | null;
      } | null
    ) => {
      validateAndSetState((prevState) => {
        const chameleonCondition = prevState.global.untimedConditions.chameleon;
        if (!chameleonCondition) {
          return prevState;
        }
        return setOneUntimedCondition(ConditionTypes.CHAMELEON, {
          ...chameleonCondition,
          targetUser: user,
        })(prevState);
      });
    },
    [validateAndSetState]
  );

  const setChameleonPersonalizeLyrics = useCallback(
    (enabled: boolean) => {
      validateAndSetState((prevState) => {
        const chameleonCondition = prevState.global.untimedConditions.chameleon;
        if (!chameleonCondition) {
          return prevState;
        }
        return setOneUntimedCondition(ConditionTypes.CHAMELEON, {
          ...chameleonCondition,
          personalizeLyrics: enabled,
        })(prevState);
      });
    },
    [validateAndSetState]
  );

  const makeSetter = useCallback(
    // this exists to create a setter with a stable reference
    // this is to avoid every call to `selectState` returning a new setter, and triggering unnecessary re-renders
    memoize(
      function <T>(path: string[]) {
        return (value: Dispatch<SetStateAction<T>>) =>
          validateAndSetState((prev) =>
            setIn(
              prev,
              path,
              typeof value === 'function' ? value(getIn(prev, path)) : value
            )
          );
      },
      (path) => path.join('.')
    ),
    [validateAndSetState]
  );

  // Create a ref to store the cache for selectState
  const selectStateCacheRef = useRef<Map<string, [unknown, unknown]>>(
    new Map()
  );

  const selectState = useCallback(
    function <T>(path: string[]): [T, Dispatch<SetStateAction<T>>] {
      const currentValue = getIn(state, path) as T;
      const setter = makeSetter(path) as Dispatch<SetStateAction<T>>;

      // Create a stable reference to the tuple
      // The setter is already memoized by makeSetter, so we only need to memoize the tuple
      const pathKey = path.join('.');

      const cached = selectStateCacheRef.current.get(pathKey);
      if (cached && cached[0] === currentValue && cached[1] === setter) {
        return cached as [T, Dispatch<SetStateAction<T>>];
      }

      const tuple: [T, Dispatch<SetStateAction<T>>] = [currentValue, setter];
      selectStateCacheRef.current.set(pathKey, tuple);
      return tuple;
    },
    [state, makeSetter]
  );

  const canClear = useMemo(() => {
    return PATHS_TO_RESET.some((path) => {
      const changed = !isEqual(
        getIn(state, path),
        getIn(defaultCreateFormState, path)
      );
      return changed;
    });
  }, [state]);

  useModelSwitchOnUpgrade(state, validateAndSetState);

  return useMemo(
    () => ({
      state,
      setState: validateAndSetState,
      selectState,

      createBlockingErrorMessage,
      canClear,
      isRemix,

      validationImposedStateChanges,
      isConfirmingValidationImposedStateChanges,
      generate,
      applyNextStateChange,

      getPromptToSave,
      lyricsPrompts,
      tagsPrompts,
      isLoadingLyricsPrompts: isLoadingLyrics,
      isLoadingTagsPrompts: isLoadingTags,
      loadLyricsPromptsNextPage: fetchNextLyricsPage,
      loadTagsPromptsNextPage: fetchNextTagsPage,
      saveLyricsPrompt,
      saveTagsPrompt,
      deletePromptFromCache,
      lastSavedLyricsPromptState,
      lastSavedTagsPromptState,

      isAddMenuOpen,
      setIsAddMenuOpen,

      addMenuState,
      setAddMenuState,

      isCreatePersonaModalOpen,
      setIsCreatePersonaModalOpen,

      isSettingsMenuOpen,
      setIsSettingsMenuOpen,

      isStudioTimelineRendering,
      setIsStudioTimelineRendering,
      lastStudioTimelineRenderTimestamp,
      setLastStudioTimelineRenderTimestamp,

      isGenerating,

      modelTypeOptions,

      creditsRemaining,
      chargeCredits,

      activeModal,
      setActiveModal,

      addPersona,
      attemptOpenPersonaModal,
      attemptOpenInspoModal,
      addCondition,
      addStudioProjectStateCondition,
      addInspoPlaylist,
      addInspirationClips,
      removeInspirationClip,
      clearInspirationClips,
      addSampleClips,
      removeSampleClip,
      clearSampleClips,

      setChameleonMode,
      getChameleonMode,
      setChameleonTargetUser,
      setChameleonPersonalizeLyrics,

      handleOverwriteConfirm,
      handleOverwriteCancel,
      handleOverwriteClose,
      handleConditionOverwriteConfirm,
      handleConditionOverwriteCancel,
      handleConditionOverwriteClose,
      pendingPersona,
      pendingCondition,
    }),
    [
      state,
      validateAndSetState,
      selectState,
      createBlockingErrorMessage,
      canClear,
      isRemix,
      validationImposedStateChanges,
      isConfirmingValidationImposedStateChanges,
      generate,
      applyNextStateChange,
      getPromptToSave,
      lyricsPrompts,
      tagsPrompts,
      isLoadingLyrics,
      isLoadingTags,
      fetchNextLyricsPage,
      fetchNextTagsPage,
      saveLyricsPrompt,
      saveTagsPrompt,
      deletePromptFromCache,
      lastSavedLyricsPromptState,
      lastSavedTagsPromptState,
      isAddMenuOpen,
      addMenuState,
      isCreatePersonaModalOpen,
      isSettingsMenuOpen,
      isStudioTimelineRendering,
      lastStudioTimelineRenderTimestamp,
      isGenerating,
      modelTypeOptions,
      creditsRemaining,
      chargeCredits,
      activeModal,
      addPersona,
      attemptOpenPersonaModal,
      attemptOpenInspoModal,
      addCondition,
      addStudioProjectStateCondition,
      addInspoPlaylist,
      addInspirationClips,
      removeInspirationClip,
      clearInspirationClips,
      addSampleClips,
      removeSampleClip,
      clearSampleClips,
      setChameleonMode,
      getChameleonMode,
      setChameleonTargetUser,
      setChameleonPersonalizeLyrics,
      handleOverwriteConfirm,
      handleOverwriteCancel,
      handleOverwriteClose,
      handleConditionOverwriteConfirm,
      handleConditionOverwriteCancel,
      handleConditionOverwriteClose,
      pendingPersona,
      pendingCondition,
    ]
  );
};

const CreateFormContext = createContext<ReturnType<typeof useCreateForm>>(
  undefined as never
);

export default CreateFormContext;