import { useAuth } from '@clerk/nextjs';
import { Allow, parse } from 'partial-json';
import { useCallback, useMemo } from 'react';
import storageAvailable from 'storage-available';

import { useStores } from '@/app/(root)/AppProviders';
import { invalidateWorkspaceQueries } from '@/components/clipBrowser/clipBrowserQueryClient';
import { ClipBrowserRegistryContext } from '@/components/clipBrowser/useClipBrowserRegistry';
import { useFetchClip } from '@/hooks/useClip';
import { useContextSelector } from '@/hooks/useContextSelector';
import useGenerate, {
  CoverReference,
  ExtendReference,
  ReferenceType as GenerateReferenceType,
  ModelTier,
  PromptType,
} from '@/hooks/useGenerate';
import { createTransactionLogger } from '@/logging/logWebUserEvent';
import logWebUserEvent from '@/logging/logWebUserEvent';
import { Clip } from '@/state/clipStore';
import { DEFAULT_CREATE_CONTROL_VALUE } from '@/utils/constants';
import { getUserLocale, isMobileBrowser } from '@/utils/device';
import { streamTextToSpeech } from '@/utils/textToSpeech';

import CreateFormContext from '../create/v2/CreateFormContext';
import resetCreateInputs from '../create/v2/actions/resetCreateInputs';
import { FocusedObjectContext } from '../create/v2/useFocusedObject';
import {
  ClipReference,
  LyricsReference,
  ReferenceType,
  StylesReference,
} from './components/input/ReferenceTypes';
import { SIMPLE_MESSAGE_TOOL_TYPES } from './constants';
import {
  MessageLikeStatus,
  useChatMessagesStore,
  useChatStore,
} from './stores';
import { useChatContext } from './useChat';
import {
  WORKSPACES_QUERY_KEY,
  workspaceListQueryClient,
} from './useWorkspaces';
import { MODAL_SRV, getUUID } from './utils';

export enum ABTestVariant {
  A = 'a',
  B = 'b',
}

export type ABTestInfo = {
  ab_test_group_id: string;
  ab_variant: string;
  selected_for_ab?: boolean;
};

export type ExecuteGenerateSongToolCallParams = {
  callId: string;
  toolArgs: any;
  isRegenerate?: boolean;
  isToolEdit?: boolean;
  batchOffset?: number;
};

const PENDING_GEN_TIMEOUT = 1000 * 8;

export const loadingTexts = [
  'Thinking...',
  'Crafting...',
  'Composing...',
  'Confabulating...',
  'Hallucinating...',
  'Noodling...',
  'Vibing...',
  'Cooking...',
  'Bopping...',
  'Pumping it up...',
  'Rizzing...',
  'Checking the vibe...',
  'Chopping it up...',
  'Discombobulating...',
];

const streamResponse = async (
  response: Response,
  handleReadChunk: (chunk: object) => boolean
) => {
  if (!response.ok) {
    console.error(`HTTP error! status: ${response.status}`);
    return false;
  }

  const reader = response.body?.getReader();
  const decoder = new TextDecoder();

  if (!reader) {
    console.error('No response body');
    return false;
  }

  try {
    while (true) {
      const { done, value } = await reader.read();

      if (done) return true;

      const chunk = decoder.decode(value, { stream: true });
      const lines = chunk.split('\n');

      for (const line of lines) {
        if (line.startsWith('data: ')) {
          const data = JSON.parse(line.slice(6));
          const ok = handleReadChunk(data);
          if (!ok) {
            reader.releaseLock();
            return false;
          }
        }
      }
    }
  } finally {
    reader.releaseLock();
  }
};

export const useChatMessageActions = () => {
  const {
    addMsgId,
    replaceMsgId,
    setMsgMapValue,
    setMsgFeedbackMapValue,
    msgIds,
    msgMap,
    msgFeedbackMap,
    setIsLoading,
    setLoadingText,
    isStreaming,
    setIsStreaming,
    setCurrentAssistantPrompt,
    resetMsgData,
    loadedMessages,
    setLoadedMessages,
  } = useChatMessagesStore();

  const toolCallPendingGensMap = useChatStore(
    (state) => state.toolCallPendingGensMap
  );
  const setToolCallPendingGensMapValue = useChatStore(
    (state) => state.setToolCallPendingGensMapValue
  );

  const rootStore = useStores();
  const { getToken } = useAuth();
  const model = useContextSelector(
    CreateFormContext,
    (context) => context.state.global.model
  );
  const setState = useContextSelector(
    CreateFormContext,
    (context) => context.setState
  );
  const generate = useGenerate();
  const clipCreated = useContextSelector(
    ClipBrowserRegistryContext,
    (context) => context.clipCreated
  );
  const fetchClip = useFetchClip();

  const {
    chatUUID,
    isChatOwner,
    workspaceId,
    references,
    clearReferences,
    ably,
    variant,
    toolArgsCache,
    lastPlayedClipId,
    addToolArgsToCache,
    setLastPlayedClipId,
    isCreatingNewSession,
    setIsCreatingNewSession,
    setWorkspaceId,
  } = useChatStore();

  const toolCallToLoadingText = (toolName: string) => {
    switch (toolName) {
      case 'create_playlist_with_clips':
        return 'Creating playlist...';
      case 'search_library':
      case 'search_public_clips_similarity':
      case 'search_public_clips_descriptive':
        return 'Searching library...';
      case 'generate_song':
        return 'Creating song...';
      case 'write_lyrics':
        return 'Writing lyrics...';
      case 'generate_image':
        return 'Creating image...';
      case 'listen_to_audio':
        return 'Listening...';
      case 'answer_music_question':
        return 'Thinking...';
      case 'publish_song':
        return 'Publishing song...';
      case 'update_song_metadata':
        return 'Updating song...';
      case 'get_download_link':
        return 'Getting download link...';
      case 'get_share_link':
        return 'Getting share link...';
      default:
        return 'Crafting...';
    }
  };

  const { attemptRegisterSession } = useChatContext();
  const focusedObject = useContextSelector(
    FocusedObjectContext,
    (context) => context?.focusedObject
  );
  const setFocusedObject = useContextSelector(
    FocusedObjectContext,
    (context) => context?.setFocusedObject
  );

  const handleAfterCreatedClips = useCallback(
    async (newClipIds: string[]) => {
      await Promise.all(
        newClipIds.map(async (clipId: string) => {
          const clip = await fetchClip(clipId, false);
          if (clip) {
            clipCreated(clip);
          }
        })
      );
      invalidateWorkspaceQueries(workspaceId ?? 'default');
      workspaceListQueryClient.invalidateQueries({
        queryKey: [WORKSPACES_QUERY_KEY],
      });
      if (focusedObject === undefined && newClipIds?.[0]) {
        setFocusedObject?.({ type: 'clip', clipId: newClipIds?.[0] });
      }
    },
    [focusedObject, setFocusedObject, clipCreated, fetchClip, workspaceId]
  );

  const registerMessageId = useCallback(
    (msgId: string, abTestInfo?: ABTestInfo) => {
      if (
        abTestInfo?.ab_test_group_id &&
        !msgIds.includes(abTestInfo?.ab_test_group_id)
      ) {
        addMsgId(abTestInfo?.ab_test_group_id);
        const variant = abTestInfo?.ab_variant.toLowerCase();
        const existingAbMapValue = msgMap.get(abTestInfo?.ab_test_group_id) ?? {
          a: [],
          b: [],
        };
        if (!(existingAbMapValue[variant] ?? []).includes(msgId)) {
          setMsgMapValue(abTestInfo?.ab_test_group_id, {
            ...existingAbMapValue,
            [variant]: [...(existingAbMapValue[variant] ?? []), msgId],
          });
        }
      } else if (
        abTestInfo?.ab_test_group_id &&
        msgIds.includes(abTestInfo?.ab_test_group_id)
      ) {
        const variant = abTestInfo?.ab_variant.toLowerCase();
        const existingAbMapValue = msgMap.get(abTestInfo?.ab_test_group_id) ?? {
          a: [],
          b: [],
        };
        if (!(existingAbMapValue[variant] ?? []).includes(msgId)) {
          setMsgMapValue(abTestInfo?.ab_test_group_id, {
            ...existingAbMapValue,
            [variant]: [...(existingAbMapValue[variant] ?? []), msgId],
          });
        }
      } else if (!msgIds.includes(msgId)) {
        addMsgId(msgId);
      }
    },
    [msgMap, msgIds]
  );

  const processToolCallDataChunk = useCallback(
    async (data: any, callId: string, name: string) => {
      const { clips } = rootStore;
      setIsLoading(false);
      const msgId = data.message_id;
      const abTestInfo = data.ab_test_info;
      // only comes back with `write_lyrics` tool responses
      const lyricsVersion = data.lyrics_version;

      registerMessageId(msgId, abTestInfo);
      if (data.type === 'content' && !!data.content) {
        if (!isStreaming) {
          setIsStreaming(true);
        }
        const mapValue: any = {
          role: 'tool',
          tool_call_id: callId,
          tool_name: name,
          message:
            (msgMap.get(msgId)?.message ?? '') +
            (typeof data.content === 'object'
              ? JSON.stringify(data.content)
              : data.content),
          message_id: msgId,
        };

        // in case the tool call response is JSON (never true for lyrics),
        // partial-parse the JSON into a separate key
        if (name !== 'write_lyrics') {
          try {
            const parsedMessage = parse(mapValue.message, Allow.ALL);
            mapValue.parsedMessage = parsedMessage;

            if (
              SIMPLE_MESSAGE_TOOL_TYPES.includes(name) &&
              !!mapValue.parsedMessage.message
            ) {
              mapValue.type = 'simple_message';
            }

            // Special handling for search_library tool call
            if (
              name === 'search_library' ||
              name === 'search_public_clips_similarity' ||
              name === 'search_public_clips_descriptive'
            ) {
              // Try different possible paths for the search results
              let searchResults = null;
              if (parsedMessage?.result?.['']?.result) {
                searchResults = parsedMessage.result[''].result;
              } else if (parsedMessage?.result) {
                searchResults = parsedMessage.result;
              }

              if (searchResults && Array.isArray(searchResults)) {
                const clipIds: string[] = [];

                // Add each clip to the store and collect clip IDs
                searchResults.forEach((clip: any) => {
                  if (clip.id) {
                    // Convert the search result to match the clip store format
                    const clipData = {
                      ...clip,
                      imageUrl: clip.image_url,
                      audioUrl: clip.audio_url,
                      videoUrl: clip.video_url,
                    };
                    clips.addClip(clipData);
                    clipIds.push(clip.id);
                  }
                });

                // Store the clip IDs for rendering with ClipsMessage
                mapValue.clipIds = clipIds;
              }
            }

            // Special handling for create_playlist_with_clips tool call
            if (name === 'create_playlist_with_clips') {
              // Extract playlist ID from the response
              let playlistId: string | null = null;

              // Try different possible paths for playlist ID
              if (parsedMessage?.playlist_id) {
                playlistId = parsedMessage.playlist_id;
              } else if (parsedMessage?.playlist?.id) {
                playlistId = parsedMessage.playlist.id;
              } else if (parsedMessage?.result?.playlist_id) {
                playlistId = parsedMessage.result.playlist_id;
              } else if (parsedMessage?.result?.playlist?.id) {
                playlistId = parsedMessage.result.playlist.id;
              }

              if (playlistId) {
                // Store the playlist ID for rendering with PlaylistMessage
                mapValue.playlistIds = [playlistId];
              }
            }
            if (
              ['generate_song', 'get_whole_song'].includes(name) &&
              parsedMessage?.clipIds
            ) {
              mapValue.clipIds = parsedMessage.clipIds;
              setLastPlayedClipId(parsedMessage.clipIds[0]);
              handleAfterCreatedClips(parsedMessage.clipIds);
            }
          } catch {}
        } else {
          mapValue.lyrics_version = lyricsVersion;
        }

        setMsgMapValue(msgId, mapValue);
      } else if (data.type === 'finished') {
        setIsStreaming(false);

        // Log lyrics generation completion
        if (name === 'write_lyrics') {
          const lyricsMessage = msgMap.get(msgId);
          if (lyricsMessage?.message) {
            logWebUserEvent(
              {
                actionName: 'OrpheusLyricsGenerated',
                context: {
                  sessionId: chatUUID,
                  toolCallId: callId,
                  lyricsLength: lyricsMessage.message.length,
                },
              },
              rootStore.session
            );
          }
        }
      }
      return true;
    },
    [
      workspaceId,
      rootStore,
      setIsLoading,
      isStreaming,
      setIsStreaming,
      addMsgId,
      setMsgMapValue,
      msgMap,
      msgIds,
      chatUUID,
    ]
  );

  const processDataChunk = useCallback(
    (data: any, isVoiceMode: boolean) => {
      const msgId = data.message_id;
      const abTestInfo = data.ab_test_info;

      registerMessageId(msgId, abTestInfo);

      switch (data.type) {
        case 'content':
          setIsLoading(false);
          if (!!data.content) {
            setMsgMapValue(msgId, (msgMap.get(msgId) ?? '') + data.content);
          }
          return true;
        case 'tool_call':
          if (msgMap.get(msgId)?.name) {
            const latestStreamedArgs =
              (msgMap.get(msgId)?.arguments ?? '') +
              data.tool_call.function.arguments;
            const latestParsedArguments =
              latestStreamedArgs.trim().length > 0
                ? parse(latestStreamedArgs, Allow.ALL)
                : {};
            const currentMsgObj = msgMap.get(msgId);
            if (
              currentMsgObj?.name === 'simple_message' &&
              (msgId === msgIds[msgIds.length - 1] ||
                abTestInfo?.ab_test_group_id === msgIds[msgIds.length - 1])
            ) {
              setIsLoading(false);
            }
            setMsgMapValue(msgId, {
              ...currentMsgObj,
              arguments: latestStreamedArgs,
              parsedArguments: latestParsedArguments,
            });
            if (latestParsedArguments?.current_prompt) {
              setCurrentAssistantPrompt(latestParsedArguments?.current_prompt);
              if (lastPlayedClipId !== null) {
                setLastPlayedClipId(null);
              }
            }
          } else {
            if (data.tool_call.function.name === 'simple_message') {
              setState(resetCreateInputs);
            } else {
              setLoadingText(
                toolCallToLoadingText(data.tool_call.function.name)
              );
            }
            setMsgMapValue(msgId, {
              type: 'tool_call',
              id: data.tool_call.id,
              ...data.tool_call.function,
            });
          }
          return true;
        case 'finished':
          const finishedMsg = msgMap.get(msgId);
          if (isVoiceMode) {
            if (finishedMsg && typeof finishedMsg === 'string') {
              try {
                const parsed = JSON.parse(finishedMsg);
                if (parsed.message) {
                  streamTextToSpeech(parsed.message);
                }
              } catch (error) {
                streamTextToSpeech(finishedMsg);
              }
            } else {
              if (finishedMsg?.arguments) {
                try {
                  const parsedArgs = JSON.parse(finishedMsg.arguments);
                  if (parsedArgs?.message) {
                    streamTextToSpeech(parsedArgs.message);
                  }
                } catch (error) {}
              }
            }
          }
          return true;
        case 'error':
          if (msgId === msgIds[msgIds.length - 1]) {
            setIsLoading(false);
          }
          return false;
        default:
          // message is sent complete from Ably, add directly to store
          if (data.message && data.message_id) {
            setMsgMapValue(data.message_id, {
              ...data,
              isCollabUserMessage: true,
              isSuggestion: data.message
                .toLowerCase()
                .split(' ')
                .includes('@orpheus'),
            });
            return true;
          }
      }
      return false;
    },
    [
      setIsLoading,
      addMsgId,
      setMsgMapValue,
      msgMap,
      msgIds,
      setState,
      setLastPlayedClipId,
      lastPlayedClipId,
      setCurrentAssistantPrompt,
    ]
  );

  const handleSend = useCallback(
    async ({
      input,
      setInput,
      message,
      forcedRole,
      toolCallId,
      clipIds,
      isHidden,
      needsResponse = true,
      isMobile = false,
      sessionIdOverride,
      // isVoiceMode,
    }: {
      input?: string;
      setInput?: (input: string) => void;
      message?: string;
      forcedRole?: string;
      toolCallId?: string;
      isVoiceMode?: boolean;
      clipIds?: string[]; // for a clip message, the explicit clip IDs
      isHidden?: boolean; // forces the message to be hidden from the chat window
      isMobile?: boolean;
      needsResponse?: boolean; // whether the message needs a response from the server
      sessionIdOverride?: string;
    }) => {
      const { session } = rootStore;

      const capturedInput = message || input || '';
      if (capturedInput.trim().length === 0) {
        return;
      }
      const captureTime = Date.now();

      // Derive legacy values from new reference system
      const clipReference = references.find(
        (ref) => ref.type === ReferenceType.CLIP
      );
      const referencedClipId = clipReference
        ? (clipReference as ClipReference).clipId
        : null;

      const lyricsReference = references.find(
        (ref) => ref.type === ReferenceType.LYRICS
      );

      const stylesReference = references.find(
        (ref) => ref.type === ReferenceType.STYLES
      );

      const newMsgId = getUUID();

      const locale = getUserLocale() ?? 'en-US';
      const platform = isMobileBrowser() ? 'mobile_web' : 'desktop_web';

      // TODO: Log a frontend event for starting the message send

      addMsgId(newMsgId);
      clearReferences();

      const newMsg: any = {
        role: forcedRole ? forcedRole : toolCallId ? 'tool' : 'user',
        message: capturedInput,
        tool_call_id: toolCallId, // snake-cased for consistency with BE
        isHiddenMessage:
          isHidden !== undefined ? isHidden : !!toolCallId && !clipIds,
        clipIds,
        referenced_clip_id: referencedClipId, // snake-cased for consistency with BE
        message_id: newMsgId,
        userId: session.userId,
      };

      if (lyricsReference) {
        newMsg.referenced_lyrics = (
          lyricsReference as LyricsReference
        ).message.slice(
          (lyricsReference as LyricsReference).selectionOverlayRange.start,
          (lyricsReference as LyricsReference).selectionOverlayRange.end
        );
        newMsg.lyrics_full_context = (
          lyricsReference as LyricsReference
        ).message;
      }

      if (stylesReference) {
        newMsg.referenced_styles = (
          stylesReference as StylesReference
        ).message.slice(
          (stylesReference as StylesReference).selectionOverlayRange?.start ??
            0,
          (stylesReference as StylesReference).selectionOverlayRange?.end ??
            (stylesReference as StylesReference).message.length
        );
        newMsg.styles_full_context = (
          stylesReference as StylesReference
        ).message;
      }

      setMsgMapValue(newMsgId, newMsg);
      const channel = ably?.channels.get(
        `orpheus-chat:${sessionIdOverride ?? chatUUID}`
      );
      if (channel) {
        channel.publish('orpheus-message', { data: newMsg });
      }
      if (!isHidden) {
        setInput?.('');
      }

      if (!isChatOwner) {
        return;
      }
      if (needsResponse) {
        setIsLoading(true);
        setLoadingText(null); // defer to IncomingMessageLoading since we don't have a contextual loading text yet
      }

      if (isCreatingNewSession) {
        setIsCreatingNewSession(false);
        attemptRegisterSession(sessionIdOverride ?? chatUUID).then(
          ({ workspaceId }) => {
            if (workspaceId) {
              setWorkspaceId(workspaceId);
            }
          }
        );
      }
      const token = await getToken();
      const payload: any = {
        message: capturedInput,
        session_id: sessionIdOverride ?? chatUUID,
        message_id: newMsgId,
        referenced_clip_id: referencedClipId,
        referenced_styles: newMsg.referenced_styles,
        styles_full_context: newMsg.styles_full_context,
        referenced_lyrics: newMsg.referenced_lyrics,
        lyrics_full_context: newMsg.lyrics_full_context,
        clip_ids: clipIds,
        store_history_only: !needsResponse,
        locale: locale,
        platform: platform,
        ab_test_probability: isMobile ? 0 : undefined,
      };
      if (toolCallId) {
        payload.tool_call_id = toolCallId;
      }
      if (forcedRole) {
        payload.role = forcedRole;
      }

      let shouldRetry = false;
      let hasExecuted = false;
      let retries = 0;
      const MAX_RETRIES = 5;
      while (retries < MAX_RETRIES && (shouldRetry || !hasExecuted)) {
        hasExecuted = true;
        shouldRetry = false;
        const timeSinceCapture = Date.now() - captureTime;
        console.log(`Message send time since capture: ${timeSinceCapture}ms`);
        console.log(payload);
        const response = await fetch(`${MODAL_SRV}/chat`, {
          method: 'POST',
          headers: {
            'Content-Type': 'application/json',
            Authorization: `Bearer ${token}`,
          },
          body: JSON.stringify({ ...payload, is_retry: retries > 0, variant }),
        });

        if (!!clipIds) {
          return;
        }
        const streamResult = await streamResponse(response, (/*data: any*/) => {
          //return processDataChunk(data);
          return true;
        });
        if (!streamResult) {
          shouldRetry = true;
          retries += 1;
        }
      }
      if (storageAvailable('localStorage') && session.userId) {
        localStorage.setItem('chat-uuid-saved-at', Date.now().toString());
        localStorage.setItem(
          `orpheus_session_last_touched:${session?.userId}`,
          Date.now().toString()
        );
        localStorage.setItem(
          `orpheus_last_workspace_id:${session?.userId}`,
          workspaceId ?? ''
        );
      }
      logWebUserEvent(
        {
          actionName: 'OrpheusMessageSent',
          context: {
            sessionId: sessionIdOverride ?? chatUUID,
            message: capturedInput,
          },
        },
        rootStore.session
      );
    },
    [
      references,
      chatUUID,
      rootStore,
      rootStore.session,
      ably,
      isChatOwner,
      getToken,
      addMsgId,
      setMsgMapValue,
      setIsLoading,
      variant,
      workspaceId,
    ]
  );

  const executeGenerateSongToolCall = useCallback(
    async ({
      callId,
      toolArgs,
      isRegenerate = false,
      isToolEdit = false,
      batchOffset,
    }: ExecuteGenerateSongToolCallParams) => {
      if (!isChatOwner) {
        console.log('Cannot execute tool call, not chat owner');
        return;
      }
      console.log('toolArgs', toolArgs);

      const { project } = rootStore;

      if (!toolArgsCache.get(callId)) {
        addToolArgsToCache(callId, toolArgs);
      }

      const transactionLogger = createTransactionLogger();
      const references = [];
      if (toolArgs.task === 'cover' && !!toolArgs.cover_clip_id) {
        references.push({
          type: GenerateReferenceType.Cover,
          clipId: toolArgs.cover_clip_id,
        } as CoverReference);
      } else if (
        toolArgs.task === 'extend' &&
        !!toolArgs.continue_clip_id &&
        !!toolArgs.continue_at
      ) {
        references.push({
          type: GenerateReferenceType.Extend,
          clipId: toolArgs.continue_clip_id,
          startSeconds: toolArgs.continue_at,
          isUpload: false,
          contextLyrics: '', // TODO fix
        } as ExtendReference);
      }

      const existingPendingGens = toolCallPendingGensMap.get(callId) || [];
      const nowTimestamp = Date.now();
      const newPendingGens = [
        ...existingPendingGens.filter(
          (genTime: number) => Date.now() - genTime < PENDING_GEN_TIMEOUT
        ),
        nowTimestamp,
      ];
      setToolCallPendingGensMapValue(callId, newPendingGens);

      try {
        const _clips = await generate(transactionLogger, {
          projectId: project.currentProjectId ?? undefined,
          prompt: {
            type: PromptType.Custom,
            lyrics: toolArgs.lyrics,
            tags: toolArgs.tags,
            title: toolArgs.title,
          },
          references: references,
          controlSliders: {
            weirdnessConstraint:
              toolArgs.weirdness_constraint !== undefined &&
              toolArgs.weirdness_constraint !== DEFAULT_CREATE_CONTROL_VALUE
                ? toolArgs.weirdness_constraint * 100.0
                : DEFAULT_CREATE_CONTROL_VALUE,
            styleWeight:
              toolArgs.style_weight !== undefined &&
              toolArgs.style_weight !== DEFAULT_CREATE_CONTROL_VALUE
                ? toolArgs.style_weight * 100.0
                : DEFAULT_CREATE_CONTROL_VALUE,
          },
          modelTier: ModelTier.V4_5,
          modelOverride: model,
          vocalGender: toolArgs.vocal_gender,
          batchOffset,
        });

        handleSend({
          message: JSON.stringify({
            clipIds: _clips.map((clip: Clip) => clip.id),
          }),
          toolCallId: callId,
          clipIds: _clips.map((clip: Clip) => clip.id),
          forcedRole: isRegenerate || isToolEdit ? 'assistant' : undefined,
          needsResponse: !isRegenerate,
        });
        invalidateWorkspaceQueries(project.currentProjectId);
      } finally {
        setToolCallPendingGensMapValue(callId, existingPendingGens);
      }
    },
    [chatUUID, getToken, rootStore, generate]
  );

  const setMessageLikeStatus = useCallback(
    async (messageId: string, likeStatus: MessageLikeStatus | null) => {
      const existingFeedback = msgFeedbackMap.get(messageId);
      setMsgFeedbackMapValue(messageId, likeStatus);
      const token = await getToken();
      const response = await fetch(`${MODAL_SRV}/message/like-status`, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          Authorization: `Bearer ${token}`,
        },
        body: JSON.stringify({
          message_id: messageId,
          like_status: likeStatus,
          session_id: chatUUID,
        }),
      });
      if (!response.ok) {
        // TODO add logging and toast message
        setMsgFeedbackMapValue(messageId, existingFeedback ?? null);
      }
    },
    [getToken, msgFeedbackMap, setMsgFeedbackMapValue, chatUUID]
  );

  const selectVariantForAB = useCallback(
    (abTestGroupId: string, variant: ABTestVariant) => {
      const msgIdsForVariant = msgMap.get(abTestGroupId)?.[variant];

      if (msgIdsForVariant) {
        replaceMsgId(abTestGroupId, msgIdsForVariant);
      } else {
        const newLoadedMessages = loadedMessages.map((loadedMessage: any) => {
          if (loadedMessage.ab_test_group_id === abTestGroupId) {
            return {
              ...loadedMessage,
              selected_for_ab: loadedMessage.ab_variant === variant,
            };
          } else {
            return { ...loadedMessage };
          }
        });
        setLoadedMessages(newLoadedMessages);
      }
    },
    [msgMap, msgIds, replaceMsgId, loadedMessages, setLoadedMessages]
  );

  return useMemo(
    () => ({
      handleSend,
      executeGenerateSongToolCall,
      resetMsgData,
      setCurrentAssistantPrompt,
      processDataChunk,
      processToolCallDataChunk,
      setMessageLikeStatus,
      selectVariantForAB,
    }),
    [
      handleSend,
      executeGenerateSongToolCall,
      resetMsgData,
      setCurrentAssistantPrompt,
      processDataChunk,
      processToolCallDataChunk,
      setMessageLikeStatus,
      selectVariantForAB,
    ]
  );
};