sa4cps-frontend/src/stores/websocket.ts

import { defineStore } from 'pinia'
import { ref, reactive } from 'vue'
import { useSensorStore } from './sensor'
import { useRoomStore } from './room'

const MAX_DATA_POINTS = 100

interface LegacyEnergyData {
  sensorId: string
  timestamp: number
  value: number
  unit: string
}

interface SensorReading {
  sensorId: string
  room: string
  timestamp: number
  energy: {
    value: number
    unit: string
  }
  co2: {
    value: number
    unit: string
  }
  temperature?: {
    value: number
    unit: string
  }
}

export const useWebSocketStore = defineStore('websocket', () => {
  // State
  const isConnected = ref(false)
  const latestMessage = ref<LegacyEnergyData | null>(null)
  const timeSeriesData = reactive<{
    labels: string[]
    datasets: { data: number[] }[]
  }>({
    labels: [],
    datasets: [{ data: [] }],
  })

  let socket: WebSocket | null = null
  const newDataBuffer: (LegacyEnergyData | SensorReading)[] = []

  // Actions
  function connect(url: string) {
    if (isConnected.value && socket) {
      console.log('Already connected.')
      return
    }

    // Close any existing connection first
    if (socket) {
      socket.onclose = null
      socket.onerror = null
      socket.onmessage = null
      socket.close()
      socket = null
    }

    console.log(`Connecting to WebSocket at ${url}`)
    socket = new WebSocket(url)

    socket.onopen = () => {
      console.log('WebSocket connection established.')
      isConnected.value = true
    }

    socket.onmessage = (event) => {
      try {
        const data = JSON.parse(event.data)

        // Handle proxy info message from API Gateway
        if (data.type === 'proxy_info' && data.sensor_service_url) {
          console.log('Received proxy info, reconnecting to sensor service...')
          // Close current connection gracefully
          if (socket) {
            socket.onclose = null // Prevent triggering disconnect handlers
            socket.close()
            socket = null
          }
          // Set disconnected state temporarily
          isConnected.value = false

          // Connect directly to sensor service after a short delay
          setTimeout(() => {
            console.log('Connecting directly to sensor service at ws://localhost:8007/ws')
            connect('ws://localhost:8007/ws')
          }, 100)
          return
        }

        newDataBuffer.push(data)
      } catch (error) {
        console.error('Error parsing incoming data:', error)
      }
    }

    socket.onclose = (event) => {
      console.log(`WebSocket connection closed. Code: ${event.code}, Reason: ${event.reason}`)
      isConnected.value = false
      socket = null
    }

    socket.onerror = (error) => {
      console.error('WebSocket error:', error)
      isConnected.value = false
      if (socket) {
        socket = null
      }
    }

    // Process the buffer at intervals
    setInterval(() => {
      if (newDataBuffer.length > 0) {
        const data = newDataBuffer.shift()
        if (data) {
          processIncomingData(data)
        }
      }
    }, 500)
  }

  function disconnect() {
    if (socket) {
      socket.close()
    }
  }

  function isLegacyData(data: any): data is LegacyEnergyData {
    return 'value' in data && !('energy' in data)
  }

  function processIncomingData(data: LegacyEnergyData | SensorReading) {
    // Skip non-data messages
    if (
      'type' in data &&
      (data.type === 'connection_established' || data.type === 'proxy_info')
    ) {
      console.log('Received system message:', data.type)
      return
    }

    const sensorStore = useSensorStore()
    const roomStore = useRoomStore()

    // Handle both legacy and new data formats
    if (isLegacyData(data)) {
      latestMessage.value = data
      sensorStore.updateSensorData(data)

      // Update time series for chart
      const newLabel = new Date(data.timestamp * 1000).toLocaleTimeString()
      timeSeriesData.labels.push(newLabel)
      timeSeriesData.datasets[0].data.push(data.value)
    } else {
      // Handle new multi-metric data
      roomStore.updateRoomData(data)

      // Update time series for chart (use energy values)
      const newLabel = new Date(data.timestamp * 1000).toLocaleTimeString()
      timeSeriesData.labels.push(newLabel)
      timeSeriesData.datasets[0].data.push(data.energy.value)
    }

    // Keep only the latest data points
    if (timeSeriesData.labels.length > MAX_DATA_POINTS) {
      timeSeriesData.labels.shift()
      timeSeriesData.datasets[0].data.shift()
    }
  }

  return {
    // State
    isConnected,
    latestMessage,
    timeSeriesData,

    // Actions
    connect,
    disconnect,
  }
})