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);
            };
        },
    };
});