An atom that allows sub attributes to be selected and empty.

[quangmuasuengx]

import { atom, SetStateAction, WritableAtom } from 'jotai';
import { has, isFunction, get as lGet } from 'lodash';

type PickProp<T, K> = K extends keyof NonNullable<T>
    ? T extends undefined
        ? NonNullable<T>[K] | undefined
        : NonNullable<T>[K]
    : undefined;

export function callSetState<T>(prev: T, setValue: SetStateAction<T>): T {
    if (isFunction(setValue)) {
        return setValue(prev);
    }
    return setValue;
}

export function focusPropAtom<Store, Result, Key extends keyof NonNullable<Store>>(
    baseAtom: WritableAtom<Store, [SetStateAction<Store>], Result>,
    path: Key,
) {
    return atom(
        (get) => {
            const value = get(baseAtom);
            return value && has(value, path) ? lGet(value, path) : undefined;
        },
        (_get, set, args) => {
            set(baseAtom, (prev) => {
                const old = (prev || {}) as unknown as NonNullable<Store>;

                const merged = {
                    ...old,
                    [path]: callSetState(old[path] as never, args),
                };

                return merged as unknown as typeof prev;
            });
        },
    ) as unknown as WritableAtom<PickProp<Store, Key>, [SetStateAction<PickProp<Store, Key>>], Result>;
}