buster/web/src/context/Chats/NewChatProvider.tsx

'use client';

import React from 'react';
import { createContext, useContextSelector } from 'use-context-selector';
import { useMemoizedFn } from '@/hooks';
import type { BusterSearchResult, FileType } from '@/api/asset_interfaces';
import { useBusterWebSocket } from '@/context/BusterWebSocket';
import { useChatStreamMessage } from './useChatStreamMessage';
import { useGetChatMemoized, useGetChatMessageMemoized } from '@/api/buster_rest/chats';
import { useChatUpdate } from './useChatUpdate';
import { create } from 'mutative';

export const useBusterNewChat = () => {
  const busterSocket = useBusterWebSocket();
  const getChatMessageMemoized = useGetChatMessageMemoized();
  const getChatMemoized = useGetChatMemoized();
  const { onUpdateChat, onUpdateChatMessage } = useChatUpdate();

  const { completeChatCallback, stopChatCallback, initializeNewChatCallback } =
    useChatStreamMessage();

  const onSelectSearchAsset = useMemoizedFn(async (asset: BusterSearchResult) => {
    await new Promise((resolve) => setTimeout(resolve, 1000));
  });

  const onStartNewChat = useMemoizedFn(
    async ({
      prompt,
      datasetId,
      metricId,
      dashboardId
    }: {
      prompt: string;
      datasetId?: string;
      metricId?: string;
      dashboardId?: string;
    }) => {
      const res = await busterSocket.emitAndOnce({
        emitEvent: {
          route: '/chats/post',
          payload: {
            dataset_id: datasetId, //TODO: add selected dataset id
            prompt,
            metric_id: metricId,
            dashboard_id: dashboardId
          }
        },
        responseEvent: {
          route: '/chats/post:initializeChat',
          callback: initializeNewChatCallback
        }
      });

      busterSocket.once({
        route: '/chats/post:complete',
        callback: completeChatCallback
      });
    }
  );

  const onStartChatFromFile = useMemoizedFn(
    async ({
      prompt,
      fileId,
      fileType
    }: {
      prompt: string;
      fileId: string;
      fileType: FileType;
    }) => {
      onStartNewChat({
        prompt,
        metricId: fileType === 'metric' ? fileId : undefined,
        dashboardId: fileType === 'dashboard' ? fileId : undefined
      });
    }
  );

  const onReplaceMessageInChat = useMemoizedFn(
    async ({
      prompt,
      messageId,
      chatId
    }: {
      prompt: string;
      messageId: string;
      chatId: string;
    }) => {
      const currentChat = getChatMemoized(chatId);
      const currentMessage = getChatMessageMemoized(messageId);
      const currentRequestMessage = currentMessage?.request_message!;
      onUpdateChatMessage({
        id: messageId,
        request_message: create(currentRequestMessage, (draft) => {
          draft.request = prompt;
        }),
        reasoning_message_ids: [],
        response_message_ids: [],
        isCompletedStream: false
      });

      const messageIndex = currentChat?.message_ids.findIndex(
        (messageId) => messageId === messageId
      );

      if (messageIndex && messageIndex !== -1) {
        const updatedMessageIds = currentChat?.message_ids.slice(0, messageIndex + 1);
        onUpdateChat({
          id: chatId,
          message_ids: updatedMessageIds
        });
      }

      await busterSocket.emitAndOnce({
        emitEvent: {
          route: '/chats/post',
          payload: {
            prompt,
            message_id: messageId,
            chat_id: chatId
          }
        },
        responseEvent: {
          route: '/chats/post:complete',
          callback: completeChatCallback
        }
      });
    }
  );

  const onFollowUpChat = useMemoizedFn(
    async ({ prompt, chatId }: { prompt: string; chatId: string }) => {
      busterSocket.once({
        route: '/chats/post:initializeChat',
        callback: initializeNewChatCallback
      });
      await busterSocket.emitAndOnce({
        emitEvent: {
          route: '/chats/post',
          payload: {
            prompt,
            chat_id: chatId
          }
        },
        responseEvent: {
          route: '/chats/post:complete',
          callback: completeChatCallback
        }
      });
    }
  );

  const onStopChat = useMemoizedFn(
    ({ chatId, messageId }: { chatId: string; messageId: string }) => {
      busterSocket.emit({
        route: '/chats/stop',
        payload: {
          id: chatId,
          message_id: messageId
        }
      });
      stopChatCallback(chatId);
    }
  );

  return {
    onStartNewChat,
    onSelectSearchAsset,
    onFollowUpChat,
    onStartChatFromFile,
    onReplaceMessageInChat,
    onStopChat
  };
};

export const BusterNewChatContext = createContext<ReturnType<typeof useBusterNewChat>>(
  {} as ReturnType<typeof useBusterNewChat>
);

export const BusterNewChatProvider: React.FC<{
  children: React.ReactNode;
}> = ({ children }) => {
  const value = useBusterNewChat();
  return <BusterNewChatContext.Provider value={value}>{children}</BusterNewChatContext.Provider>;
};

export const useBusterNewChatContextSelector = <T,>(
  selector: (state: ReturnType<typeof useBusterNewChat>) => T
) => useContextSelector(BusterNewChatContext, selector);