web-map/src/hooks/usePlaybackWebSocket.ts

import { message } from 'ant-design-vue';
import { type Ref, ref, type ShallowRef, shallowRef } from 'vue';

import type { RobotRealtimeInfo } from '../apis/robot';
import type { EditorService } from '../services/editor.service';
import ws from '../services/ws';

// Define the structure of WebSocket messages for playback
type PlaybackMessage = {
  type: 'AMR' | 'SCENE';
  timestamp: number;
  data: any;
};

// Hook's return type definition
export interface UsePlaybackWebSocketReturn {
  // Connection status
  isConnected: Ref<boolean>;

  // Playback state
  isPlaying: Ref<boolean>;
  currentTime: Ref<number>;
  totalDuration: Ref<number>;

  // Data for rendering
  currentSceneId: Ref<string | null>;
  sceneJson: Ref<string | null>;

  // Methods
  connect: (historySceneId: string) => void;
  disconnect: () => void;
  play: () => void;
  pause: () => void;
  seek: (timestamp: number) => void;
  changeSpeed: (speed: number) => void; // Placeholder for speed control
}

export function usePlaybackWebSocket(editorService: ShallowRef<EditorService | undefined>): UsePlaybackWebSocketReturn {
  const client = shallowRef<WebSocket | null>(null);
  const isConnected = ref(false);
  const isPlaying = ref(false);

  const currentTime = ref(0);
  const totalDuration = ref(0);
  const currentSceneId = ref<string | null>(null);
  const sceneJson = ref<string | null>(null);

  const latestRobotData = new Map<string, RobotRealtimeInfo>();
  let animationFrameId: number;

  const connect = async (historySceneId: string) => {
    if (client.value) {
      disconnect();
    }

    const wsPath = `/scene/monitor/logplayback/${historySceneId}`;

    try {
      const wsInstance = await ws.create(wsPath, {
        heartbeatInterval: 3600000, // 1 hour
      });
      client.value = wsInstance;

      wsInstance.onopen = () => {
        isConnected.value = true;
        message.success('回放连接已建立');
        startRenderLoop();
      };

      wsInstance.onmessage = (event) => {
        const msg = JSON.parse(event.data) as PlaybackMessage;

        if (msg.timestamp) {
          currentTime.value = msg.timestamp;
        }

        if (msg.type === 'SCENE') {
          sceneJson.value = msg.data;
          latestRobotData.clear();
        } else if (msg.type === 'AMR') {
          (msg.data as RobotRealtimeInfo[]).forEach(robot => {
            latestRobotData.set(robot.id, robot);
          });
        }
      };

      wsInstance.onclose = () => {
        isConnected.value = false;
        isPlaying.value = false;
        stopRenderLoop();
        message.info('回放连接已断开');
        client.value = null; // Clean up the client ref
      };

      wsInstance.onerror = (error) => {
        console.error('回放 WebSocket 发生错误:', error);
        message.error('回放连接发生错误');
        disconnect();
      };
    } catch (error) {
      console.error('创建 WebSocket 连接失败:', error);
      message.error('创建回放连接失败');
    }
  };

  const disconnect = () => {
    if (client.value) {
      client.value.close();
      client.value = null;
    }
  };

  const sendCommand = (command: string) => {
    if (client.value && client.value.readyState === WebSocket.OPEN) {
      client.value.send(command);
    } else {
      message.warn('回放连接未建立或已断开，无法发送指令');
    }
  };

  const play = () => {
    sendCommand('PLAY');
    isPlaying.value = true;
  };

  const pause = () => {
    sendCommand(`PAUSE:${currentTime.value}`);
    isPlaying.value = false;
  };

  const seek = (timestamp: number) => {
    sendCommand(`SEEK:${timestamp}`);
  };

  const changeSpeed = (speed: number) => {
    console.warn(`播放速度控制 (${speed}x) 尚未实现`);
  };

  const batchUpdateRobots = (updates: Array<{ id: string; data: RobotRealtimeInfo }>) => {
    const editor = editorService.value;
    if (!editor || updates.length === 0) return;

    updates.forEach(({ id, data }) => {
      if (editor.checkRobotById(id)) {
        const { x, y, angle, ...rest } = data;
        editor.updateRobot(id, rest);
        editor.setValue({
          id,
          x: x - 60,
          y: y - 60,
          rotate: -angle! + 180,
          visible: true,
        }, { render: false });
      }
    });
  };

  const renderLoop = () => {
    const updates: Array<{ id: string; data: RobotRealtimeInfo }> = [];
    for (const [id, data] of latestRobotData.entries()) {
      updates.push({ id, data });
    }
    latestRobotData.clear();

    if (updates.length > 0) {
      batchUpdateRobots(updates);
    }

    // if (isPlaying.value && currentTime.value >= totalDuration.value && totalDuration.value > 0) {
    //   pause();
    //   currentTime.value = totalDuration.value;
    // }

    editorService.value?.render();
    animationFrameId = requestAnimationFrame(renderLoop);
  };

  const startRenderLoop = () => {
    stopRenderLoop();
    renderLoop();
  };

  const stopRenderLoop = () => {
    cancelAnimationFrame(animationFrameId);
  };

  return {
    isConnected,
    isPlaying,
    currentTime,
    totalDuration,
    currentSceneId,
    sceneJson,
    connect,
    disconnect,
    play,
    pause,
    seek,
    changeSpeed,
  };
}