Cyrene/frontend/web/src/hooks/useWebSocket.ts

import { useEffect, useRef, useCallback, useState } from 'react';
import { useChatStore } from '@/store/chatStore';
import { useSessionStore } from '@/store/sessionStore';
import { getToken } from '@/api/client';
import type { WSClientMessage, WSServerMessage } from '@/types/chat';

const WS_BASE_URL =
  import.meta.env.VITE_WS_URL || 'ws://localhost:8080/ws/chat';

export function useWebSocket() {
  const [isConnected, setIsConnected] = useState(false);
  const wsRef = useRef<WebSocket | null>(null);
  const reconnectTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
  const shouldReconnectRef = useRef(true);
  const activeSessionRef = useRef<string | null>(null);

  // 订阅 sessionStore 中的 currentSessionId 变化
  const currentSessionId = useSessionStore((s) => s.currentSessionId);

  const connect = useCallback(() => {
    const token = getToken();
    if (!token) return;

    // 关闭旧连接
    if (wsRef.current) {
      shouldReconnectRef.current = false;
      wsRef.current.close();
      wsRef.current = null;
    }

    const sessionID = useSessionStore.getState().currentSessionId || '';
    const url = sessionID
      ? `${WS_BASE_URL}?token=${token}&session_id=${sessionID}`
      : `${WS_BASE_URL}?token=${token}`;

    const ws = new WebSocket(url);

    ws.onopen = () => {
      setIsConnected(true);
      shouldReconnectRef.current = true;
      console.log('[WS] 已连接, session_id:', sessionID);

      // 连接后发送会话恢复消息，恢复后端上下文
      const sid = useSessionStore.getState().currentSessionId;
      if (sid) {
        const resumeMsg: WSClientMessage = {
          type: 'history',
          session_id: sid,
          timestamp: Date.now(),
        };
        ws.send(JSON.stringify(resumeMsg));
        console.log('[WS] 已发送会话恢复请求, session_id:', sid);
      }
    };

    ws.onclose = () => {
      setIsConnected(false);
      console.log('[WS] 已断开');
      if (shouldReconnectRef.current) {
        console.log('[WS] 3秒后重连...');
        if (reconnectTimerRef.current) {
          clearTimeout(reconnectTimerRef.current);
        }
        reconnectTimerRef.current = setTimeout(() => connect(), 3000);
      }
    };

    ws.onerror = (err) => {
      console.error('[WS] 连接错误:', err);
    };

    ws.onmessage = (event) => {
      try {
        const msg: WSServerMessage = JSON.parse(event.data);
        handleServerMessage(msg);
      } catch (err) {
        console.error('[WS] 消息解析失败:', err);
      }
    };

    wsRef.current = ws;
  }, []);

  // 会话切换时：重建 WebSocket 连接（消息历史由 useSession.setCurrentSession 负责加载）
  useEffect(() => {
    activeSessionRef.current = currentSessionId;

    connect();

    return () => {
      if (reconnectTimerRef.current) {
        clearTimeout(reconnectTimerRef.current);
      }
      shouldReconnectRef.current = false;
      wsRef.current?.close();
    };
  }, [connect, currentSessionId]);

  const sendMessage = useCallback((msg: WSClientMessage) => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      const sessionID = useSessionStore.getState().currentSessionId;
      wsRef.current.send(
        JSON.stringify({
          ...msg,
          session_id: msg.session_id || sessionID || undefined,
        })
      );
    }
  }, []);

  return { isConnected, sendMessage };
}

function handleServerMessage(msg: WSServerMessage) {
  const { addMessage, appendToLastMessage, finishStreaming, setTyping } =
    useChatStore.getState();
  const { setMessages } = useSessionStore.getState();
  const chatState = useChatStore.getState();

  switch (msg.type) {
    case 'response':
      if (msg.text) {
        addMessage({
          id: msg.message_id || '',
          role: 'assistant',
          content: msg.text,
          timestamp: msg.timestamp,
        });
      }
      setTyping(false);
      break;

    case 'stream_chunk':
      if (msg.content) {
        const { messages } = useChatStore.getState();
        const lastMsg = messages[messages.length - 1];
        if (
          !lastMsg ||
          lastMsg.role !== 'assistant' ||
          !lastMsg.isStreaming
        ) {
          addMessage({
            id: msg.message_id || 'msg_' + Date.now(),
            role: 'assistant',
            content: msg.content,
            timestamp: msg.timestamp,
            isStreaming: true,
          });
          setTyping(false);
        } else {
          appendToLastMessage(msg.content);
        }
      }
      break;

    case 'stream_end':
      finishStreaming();
      break;

    case 'history_response':
      if (msg.messages) {
        const msgsWithIds = msg.messages.map((m: any, i: number) => ({
          ...m,
          id: m.id || `hist_${i}_${Date.now()}`,
        }));
        setMessages(msgsWithIds);
        useChatStore.getState().setMessages(msgsWithIds);
      }
      setTyping(false);
      break;

    case 'device_update':
      if (msg.devices && msg.devices.length > 0) {
        chatState.setIoTDevices(msg.devices);
      }
      break;

    case 'background_thinking':
      if (msg.thinking_status) {
        chatState.setBackgroundThinkingStatus(msg.thinking_status);
      }
      break;

    case 'error':
      console.error('[WS] 服务端错误:', msg.error);
      setTyping(false);
      break;

    case 'pong':
      break;

    default:
      break;
  }
}