memowake-front/hooks/useWebSocketSubscription.ts

import { useCallback, useEffect, useRef } from 'react';
import { useTranslation } from 'react-i18next';
import { getWebSocketManager, WsMessage, getWebSocketErrorMessage } from '@/lib/websocket-util';
import { Message, Assistant } from '@/types/ask';
import { Dispatch, SetStateAction } from 'react';

export const useWebSocketSubscription = (
  setUserMessages: Dispatch<SetStateAction<Message[]>>,
  isMounted: boolean
) => {
  const { t } = useTranslation();
  const isMountedRef = useRef(isMounted);
  
  // 更新挂载状态
  useEffect(() => {
    isMountedRef.current = isMounted;
  }, [isMounted]);

  const handleChatStream = useCallback((message: WsMessage) => {
    // 确保组件仍然挂载
    if (!isMountedRef.current) return;

    if (message.type === 'ChatStream') {
      setUserMessages(prevMessages => {
        // 使用 try-catch 包装以防止错误导致应用崩溃
        try {
          const lastMessage = prevMessages[prevMessages.length - 1];

          // 检查是否是有效的助手消息
          if (!lastMessage || lastMessage.role !== Assistant) {
            return prevMessages;
          }

          // 创建新的消息数组
          const newMessages = [...prevMessages];

          if (typeof lastMessage.content === 'string') {
            if (lastMessage.content === 'keepSearchIng') {
              // 第一次收到流式消息，替换占位符
              newMessages[newMessages.length - 1] = {
                ...lastMessage,
                content: message.chunk
              };
            } else {
              // 持续追加流式消息
              newMessages[newMessages.length - 1] = {
                ...lastMessage,
                content: lastMessage.content + message.chunk
              };
            }
          } else if (Array.isArray(lastMessage.content)) {
            // 如果 content 是数组，则更新第一个 text 部分
            const textPartIndex = lastMessage.content.findIndex(p => p.type === 'text');
            if (textPartIndex !== -1) {
              const updatedContent = [...lastMessage.content];
              updatedContent[textPartIndex] = {
                ...updatedContent[textPartIndex],
                text: (updatedContent[textPartIndex].text || '') + message.chunk
              };

              newMessages[newMessages.length - 1] = {
                ...lastMessage,
                content: updatedContent
              };
            }
          }

          return newMessages;
        } catch (error) {
          console.error('处理 ChatStream 消息时出错:', error);
          // 发生错误时返回原始消息数组
          return prevMessages;
        }
      });
    }
  }, [setUserMessages, isMounted]);

  const handleChatStreamEnd = useCallback((message: WsMessage) => {
    // 确保组件仍然挂载
    if (!isMountedRef.current) return;

    if (message.type === 'ChatStreamEnd') {
      setUserMessages(prevMessages => {
        // 使用 try-catch 包装以防止错误导致应用崩溃
        try {
          const lastMessage = prevMessages[prevMessages.length - 1];

          // 检查是否是有效的助手消息
          if (!lastMessage || lastMessage.role !== Assistant) {
            return prevMessages;
          }

          // 使用最终消息替换流式消息，确保 message.message 存在
          if (message.message) {
            const newMessages = [...prevMessages];
            newMessages[newMessages.length - 1] = message.message as Message;
            return newMessages;
          } else {
            // 如果最终消息为空，则移除 'keepSearchIng' 占位符
            return prevMessages.filter(m =>
              !(typeof m.content === 'string' && m.content === 'keepSearchIng')
            );
          }
        } catch (error) {
          console.error('处理 ChatStreamEnd 消息时出错:', error);
          // 发生错误时返回原始消息数组
          return prevMessages;
        }
      });
    }
  }, [setUserMessages, isMounted]);

  const handleError = useCallback((message: WsMessage) => {
    // 确保组件仍然挂载
    if (!isMountedRef.current) return;

    if (message.type === 'Error') {
      console.log(`WebSocket Error: ${message.code} - ${message.message}`);

      setUserMessages(prevMessages => {
        // 使用 try-catch 包装以防止错误导致应用崩溃
        try {
          const lastMessage = prevMessages[prevMessages.length - 1];

          // 检查是否是有效的助手消息且包含占位符
          if (!lastMessage ||
            lastMessage.role !== Assistant ||
            typeof lastMessage.content !== 'string' ||
            lastMessage.content !== 'keepSearchIng') {
            return prevMessages;
          }

          // 替换占位符为错误消息
          const newMessages = [...prevMessages];
          newMessages[newMessages.length - 1] = {
            ...lastMessage,
            content: getWebSocketErrorMessage(message.code, t)
          };

          return newMessages;
        } catch (error) {
          console.error('处理 Error 消息时出错:', error);
          // 发生错误时返回原始消息数组
          return prevMessages;
        }
      });
    }
  }, [setUserMessages, t]);

  const subscribeToWebSocket = useCallback(() => {
    const webSocketManager = getWebSocketManager();
    webSocketManager.connect();

    webSocketManager.subscribe('ChatStream', handleChatStream);
    webSocketManager.subscribe('ChatStreamEnd', handleChatStreamEnd);
    webSocketManager.subscribe('Error', handleError);

    return () => {
      // 清理订阅
      webSocketManager.unsubscribe('ChatStream', handleChatStream);
      webSocketManager.unsubscribe('ChatStreamEnd', handleChatStreamEnd);
      webSocketManager.unsubscribe('Error', handleError);

      // 不要在这里断开连接，因为其他组件可能还在使用
      // webSocketManager.disconnect();
    };
  }, [handleChatStream, handleChatStreamEnd, handleError]);

  return {
    subscribeToWebSocket,
    handleChatStream,
    handleChatStreamEnd,
    handleError
  };
};