OptimCloud-gw-ui/src/contexts/SecuritySocketProvider/useStore.ts

import { QueryClient } from '@tanstack/react-query';
import { v4 as uuid } from 'uuid';
import create from 'zustand';
import { SecuritySocketRawMessage, SocketEventCallback, SocketWebSocketNotificationData } from './utils';
import { axiosSec } from 'constants/axiosInstances';
import { NotificationType } from 'models/Socket';

export type SecurityWebSocketMessage =
  | {
      type: 'NOTIFICATION';
      data: SocketWebSocketNotificationData;
      timestamp: Date;
      id: string;
    }
  | {
      type: 'UNKNOWN';
      data: {
        type?: string;
        type_id?: number;
        [key: string]: unknown;
      };
      timestamp: Date;
      id: string;
    };

const parseRawWebSocketMessage = (message?: SecuritySocketRawMessage): SocketWebSocketNotificationData | undefined => {
  if (message && message.notification) {
    if (message.notification.type_id === 1) {
      return {
        type: 'LOG',
        log: message.notification.content,
      };
    }
  } else if (message?.notificationTypes) {
    return {
      type: 'INITIAL_MESSAGE',
      message,
    };
  }
  return undefined;
};

export type SecurityStoreState = {
  availableLogTypes: NotificationType[];
  hiddenLogIds: number[];
  setHiddenLogIds: (logsToHide: number[]) => void;
  lastMessage?: SecurityWebSocketMessage;
  allMessages: SecurityWebSocketMessage[];
  addMessage: (rawMsg: SecuritySocketRawMessage, queryClient: QueryClient) => void;
  eventListeners: SocketEventCallback[];
  addEventListeners: (callback: SocketEventCallback[]) => void;
  webSocket?: WebSocket;
  send: (str: string) => void;
  startWebSocket: (token: string, tries?: number) => void;
  isWebSocketOpen: boolean;
  setWebSocketOpen: (isOpen: boolean) => void;
  lastSearchResults: string[];
  setLastSearchResults: (result: string[]) => void;
  errors: { str: string; timestamp: Date }[];
};

export const useSecurityStore = create<SecurityStoreState>((set, get) => ({
  availableLogTypes: [],
  hiddenLogIds: [],
  setHiddenLogIds: (logsToHide: number[]) => {
    get().send(JSON.stringify({ 'drop-notifications': logsToHide }));
    set(() => ({
      hiddenLogIds: logsToHide,
    }));
  },
  allMessages: [] as SecurityWebSocketMessage[],
  addMessage: (rawMsg: SecuritySocketRawMessage) => {
    try {
      const msg = parseRawWebSocketMessage(rawMsg);
      if (msg) {
        // Handle notification-specific logic
        if (msg.type === 'INITIAL_MESSAGE') {
          if (msg.message.notificationTypes) {
            set({ availableLogTypes: msg.message.notificationTypes });
          }
        }
        // General handling
        const obj: SecurityWebSocketMessage = {
          type: 'NOTIFICATION',
          data: msg,
          timestamp: msg.log?.timestamp ? new Date(msg.log.timestamp * 1000) : new Date(),
          id: uuid(),
        };

        const eventsToFire = get().eventListeners.filter(
          ({ type, serialNumber }) => type === msg.type && serialNumber === msg.serialNumber,
        );

        if (eventsToFire.length > 0) {
          for (const event of eventsToFire) {
            event.callback();
          }

          return set((state) => ({
            allMessages:
              state.allMessages.length <= 1000 ? [...state.allMessages, obj] : [...state.allMessages.slice(1), obj],
            lastMessage: obj,
            eventListeners: get().eventListeners.filter(
              ({ id }) => !eventsToFire.find(({ id: findId }) => findId === id),
            ),
          }));
        }

        return set((state) => ({
          allMessages:
            state.allMessages.length <= 1000 ? [...state.allMessages, obj] : [...state.allMessages.slice(1), obj],
          lastMessage: obj,
        }));
      }
      return undefined;
    } catch {
      // TODO - Add error message to socket logs
      return set((state) => ({
        errors: [...state.errors, { str: JSON.stringify(rawMsg), timestamp: new Date() }],
      }));
    }
  },
  eventListeners: [] as SocketEventCallback[],
  addEventListeners: (events: SocketEventCallback[]) =>
    set((state) => ({ eventListeners: [...state.eventListeners, ...events] })),
  isWebSocketOpen: false,
  setWebSocketOpen: (isOpen: boolean) => set({ isWebSocketOpen: isOpen }),
  send: (str: string) => {
    const ws = get().webSocket;
    if (ws) ws.send(str);
  },
  startWebSocket: (token: string, tries = 0) => {
    const newTries = tries + 1;
    if (tries <= 10) {
      set({
        webSocket: new WebSocket(
          `${
            axiosSec?.defaults?.baseURL ? axiosSec.defaults.baseURL.replace('https', 'wss').replace('http', 'ws') : ''
          }/ws`,
        ),
      });
      const ws = get().webSocket;
      if (ws) {
        ws.onopen = () => {
          set({ isWebSocketOpen: true });
          ws.send(`token:${token}`);
        };
        ws.onclose = () => {
          set({ isWebSocketOpen: false });
          setTimeout(() => get().startWebSocket(token, newTries), 3000);
        };
      }
    }
  },
  lastSearchResults: [] as string[],
  setLastSearchResults: (results: string[]) => set({ lastSearchResults: results }),
  errors: [],
}));