import type { StateCreator, StoreMutatorIdentifier } from '../vanilla'

type SubscribeWithSelector = <
    T,
    Mps extends [StoreMutatorIdentifier, unknown][] = [],
    Mcs extends [StoreMutatorIdentifier, unknown][] = []
>(
    initializer: StateCreator<T, [...Mps, ['zustand/subscribeWithSelector', never]], Mcs>
) => StateCreator<T, Mps, [['zustand/subscribeWithSelector', never], ...Mcs]>

type Write<T, U> = Omit<T, keyof U> & U

type WithSelectorSubscribe<S> = S extends { getState: () => infer T } ? Write<S, StoreSubscribeWithSelector<T>> : never

declare module '../vanilla' {
    // eslint-disable-next-line @typescript-eslint/no-unused-vars
    interface StoreMutators<S, A> {
        ['zustand/subscribeWithSelector']: WithSelectorSubscribe<S>
    }
}

type StoreSubscribeWithSelector<T> = {
    subscribe: {
        (listener: (selectedState: T, previousSelectedState: T) => void): () => void
        <U>(
            selector: (state: T) => U,
            listener: (selectedState: U, previousSelectedState: U) => void,
            options?: {
                equalityFn?: (a: U, b: U) => boolean
                fireImmediately?: boolean
            }
        ): () => void
    }
}

type SubscribeWithSelectorImpl = <T extends object>(
    storeInitializer: StateCreator<T, [], []>
) => StateCreator<T, [], []>

const subscribeWithSelectorImpl: SubscribeWithSelectorImpl = (fn) => (set, get, api) => {
    type S = ReturnType<typeof fn>
    type Listener = (state: S, previousState: S) => void
    const origSubscribe = api.subscribe as (listener: Listener) => () => void
    api.subscribe = ((selector: any, optListener: any, options: any) => {
        let listener: Listener = selector // if no selector
        if (optListener) {
            const equalityFn = options?.equalityFn || Object.is
            let currentSlice = selector(api.getState())
            listener = (state) => {
                const nextSlice = selector(state)
                if (!equalityFn(currentSlice, nextSlice)) {
                    const previousSlice = currentSlice
                    optListener((currentSlice = nextSlice), previousSlice)
                }
            }
            if (options?.fireImmediately) {
                optListener(currentSlice, currentSlice)
            }
        }
        return origSubscribe(listener)
    }) as any
    const initialState = fn(set, get, api)
    return initialState
}
export const subscribeWithSelector = subscribeWithSelectorImpl as unknown as SubscribeWithSelector