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

    [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 () => {
                console.log("removing voice server update handler", handler);
                voiceHandlers.delete(handler);
            };
        }
    };
});