import { useQuery } from "@tanstack/react-query"
import { create as batshitCreate, keyResolver } from "@yornaath/batshit"
import { create } from "zustand"
import { immer } from "zustand/middleware/immer"
import { getUser } from "~/lib/api/client/user"
import type { FullUser, PartialUser, UserId } from "~/lib/api/types"

type UsersStore = {
    users: Record<UserId, PartialUser>
    currentUserId: UserId | undefined
    fetchUsersIfNotPresent: (userIds: UserId[]) => Promise<void>
    addUser: (user: PartialUser) => void
    removeUser: (userId: UserId) => void
    setCurrentUserId: (userId: UserId) => void
    getCurrentUser: () => FullUser | undefined
}

const usersFetcher = batshitCreate({
    fetcher: async (userIds: UserId[]) => {
        let users = []

        for (const userId of userIds) {
            users.push(getUser(userId))
        }

        return await Promise.all(users)
    },
    resolver: keyResolver("id")
})

export const useUserQuery = (userId: UserId) => useQuery(
    {
        queryKey: ["users", userId],
        queryFn: async () => {
            const user = await getUser(userId)
            return user
        },
        select: (data) => {
            useUsersStore.getState().addUser(data)
            return data
        }
    }
)

export const useUsersStore = create<UsersStore>()(
    immer(
        (set, get) => ({
            users: {},
            currentUserId: undefined,
            fetchUsersIfNotPresent: async (userIds) => {
                let userPromises: Promise<PartialUser>[] = []
                for (const userId of userIds) {
                    const user = get().users[userId]
                    if (!user) {
                        userPromises.push(usersFetcher.fetch(userId))
                    }
                }

                const users = await Promise.all(userPromises)
                const activeUsers = users.filter(Boolean)

                set((state) => {
                    for (const user of activeUsers) {
                        if (user?.id)
                            state.users[user.id] = user
                    }

                })
            },
            addUser: (user) => set((state) => {
                if (user.id !== get().currentUserId)
                    state.users[user.id] = user
                else {
                    const currentUser = get().users[user.id]
                    if (currentUser)
                        state.users[user.id] = { ...currentUser, ...user }
                    else
                        state.users[user.id] = user
                }
            }),
            removeUser: (userId) => set((state) => {
                delete state.users[userId]
            }),

            setCurrentUserId: (userId) => set((state) => {
                state.currentUserId = userId
            }),

            getCurrentUser: () => !!get().currentUserId ? get().users[get().currentUserId!] as FullUser : undefined
        }),
    )
)