diplom-frontend/app/stores/gateway-store.ts

import type { QueryClient } from "@tanstack/react-query";
import { create } from "zustand";
import { messageSchema, type ChannelId, type Message, type MessageId, type ServerId } from "~/lib/api/types";
import { GatewayClient } from "~/lib/websocket/gateway/client";
import { ConnectionState, EventType, type EventData, type VoiceServerUpdateEvent } from "~/lib/websocket/gateway/types";
import { useChannelsVoiceStateStore } from "./channels-voice-state";
import { usePrivateChannelsStore } from "./private-channels-store";
import { useServerChannelsStore } from "./server-channels-store";
import { useServerListStore } from "./server-list-store";
import { useUsersStore } from "./users-store";

const GATEWAY_URL = "ws://localhost:12345/gateway/ws";

const HANDLERS = {
    [EventType.ADD_SERVER]: (self: GatewayState, data: Extract<EventData, { type: EventType.ADD_SERVER }>["data"]) => {
        useServerListStore.getState().addServer(data.server);
    },

    [EventType.REMOVE_SERVER]: (
        self: GatewayState,
        data: Extract<EventData, { type: EventType.REMOVE_SERVER }>["data"],
    ) => {
        useServerListStore.getState().removeServer(data.serverId);
        useServerChannelsStore.getState().removeServer(data.serverId);
        useChannelsVoiceStateStore.getState().removeChannel(data.serverId);
    },

    [EventType.ADD_DM_CHANNEL]: (
        self: GatewayState,
        data: Extract<EventData, { type: EventType.ADD_DM_CHANNEL }>["data"],
    ) => {
        usePrivateChannelsStore.getState().addChannel({
            ...data.channel,
            recipients: data.recipients,
        });
    },

    [EventType.REMOVE_DM_CHANNEL]: (
        self: GatewayState,
        data: Extract<EventData, { type: EventType.REMOVE_DM_CHANNEL }>["data"],
    ) => {
        usePrivateChannelsStore.getState().removeChannel(data.channelId);
        useChannelsVoiceStateStore.getState().removeChannel(data.channelId);
    },

    [EventType.ADD_SERVER_CHANNEL]: (
        self: GatewayState,
        data: Extract<EventData, { type: EventType.ADD_SERVER_CHANNEL }>["data"],
    ) => {
        useServerChannelsStore.getState().addChannel(data.channel);
    },

    [EventType.REMOVE_SERVER_CHANNEL]: (
        self: GatewayState,
        data: Extract<EventData, { type: EventType.REMOVE_SERVER_CHANNEL }>["data"],
    ) => {
        useServerChannelsStore.getState().removeChannel(data.serverId, data.channelId);
        useChannelsVoiceStateStore.getState().removeChannel(data.serverId);
    },

    [EventType.ADD_USER]: (self: GatewayState, data: Extract<EventData, { type: EventType.ADD_USER }>["data"]) => {
        useUsersStore.getState().addUser(data.user);
    },

    [EventType.REMOVE_USER]: (
        self: GatewayState,
        data: Extract<EventData, { type: EventType.REMOVE_USER }>["data"],
    ) => {
        useUsersStore.getState().removeUser(data.userId);
    },

    [EventType.ADD_SERVER_MEMBER]: (
        self: GatewayState,
        data: Extract<EventData, { type: EventType.ADD_SERVER_MEMBER }>["data"],
    ) => {
        useUsersStore.getState().addUser(data.user);
    },

    [EventType.REMOVE_SERVER_MEMBER]: (
        self: GatewayState,
        data: Extract<EventData, { type: EventType.REMOVE_SERVER_MEMBER }>["data"],
    ) => {
        useUsersStore.getState().removeUser(data.userId);
    },

    [EventType.ADD_MESSAGE]: (
        self: GatewayState,
        data: Extract<EventData, { type: EventType.ADD_MESSAGE }>["data"],
    ) => {
        const message = messageSchema.parse(data.message);

        if (self.queryClient) {
            self.queryClient.setQueryData(
                ["messages", message.channelId],
                (oldData: { pages: Message[][]; pageParams: MessageId[] }) => {
                    return {
                        pages: oldData?.pages
                            ? [[message, ...oldData.pages[0]], ...oldData.pages.slice(1)]
                            : [[message]],
                        pageParams: oldData?.pageParams ?? [undefined, message.id],
                    };
                },
            );
        }
    },

    [EventType.REMOVE_MESSAGE]: (
        self: GatewayState,
        data: Extract<EventData, { type: EventType.REMOVE_MESSAGE }>["data"],
    ) => {
        if (self.queryClient) {
            self.queryClient.setQueryData(["messages", data.channelId], (oldData: any) => {
                if (!oldData) return [];
                return oldData.filter((message: any) => message.id !== data.messageId);
            });
        }
    },

    [EventType.VOICE_CHANNEL_CONNECTED]: (
        self: GatewayState,
        data: Extract<EventData, { type: EventType.VOICE_CHANNEL_CONNECTED }>["data"],
    ) => {
        useChannelsVoiceStateStore.getState().addUser(data.channelId, data.userId, {
            deaf: false,
            muted: false,
        });
    },

    [EventType.VOICE_CHANNEL_DISCONNECTED]: (
        self: GatewayState,
        data: Extract<EventData, { type: EventType.VOICE_CHANNEL_DISCONNECTED }>["data"],
    ) => {
        useChannelsVoiceStateStore.getState().removeUser(data.channelId, data.userId);
    },
};

interface GatewayState {
    client: GatewayClient | null;
    queryClient: QueryClient | null;
    status: ConnectionState;

    connect: (token: string) => void;
    disconnect: () => void;

    setQueryClient: (client: QueryClient) => void;

    updateVoiceState: (serverId: ServerId, channelId: ChannelId) => void;
    requestVoiceStates: (serverId: ServerId) => void;
    onVoiceServerUpdate: (handler: (event: VoiceServerUpdateEvent["data"]) => void | Promise<void>) => () => void;
}

export const useGatewayStore = create<GatewayState>()((set, get) => {
    const client = new GatewayClient(GATEWAY_URL);

    const voiceHandlers = new Set<(event: VoiceServerUpdateEvent["data"]) => void>();

    client.onEvent(EventType.VOICE_SERVER_UPDATE, (event) => {
        voiceHandlers.forEach((handler) => handler(event));
    });

    for (const [type, handler] of Object.entries(HANDLERS)) {
        client.onEvent(type, (data: any) => {
            handler(get(), data);
        });
    }

    return {
        client,
        queryClient: null,
        status: ConnectionState.DISCONNECTED,

        connect: (token) => {
            client.connect(token);
            set({ status: client.connectionState });

            client.onControl("stateChange", (state) => {
                set({ status: state });
            });
        },

        disconnect: () => {
            client.disconnect();
            set({ status: ConnectionState.DISCONNECTED });
        },

        setQueryClient: (queryClient) => {
            set({ queryClient });
        },

        updateVoiceState: (serverId, channelId) => {
            client.updateVoiceState(serverId, channelId);
        },

        requestVoiceStates: (serverId) => {
            client.requestVoiceStates(serverId);
        },

        onVoiceServerUpdate: (handler) => {
            voiceHandlers.add(handler);

            return () => {
                voiceHandlers.delete(handler);
            };
        },
    };
});