updux/src/Updux.ts

import * as R from 'remeda';
import {
    createStore as reduxCreateStore,
    applyMiddleware,
    DeepPartial,
    Action,
    MiddlewareAPI,
    AnyAction,
    Middleware,
    Dispatch,
} from 'redux';
import {
    configureStore,
    Reducer,
    ActionCreator,
    ActionCreatorWithoutPayload,
    ActionCreatorWithPreparedPayload,
} from '@reduxjs/toolkit';
import { AggregateActions, AggregateSelectors, Dux } from './types.js';
import { buildActions } from './buildActions.js';
import { buildInitial, AggregateState } from './initial.js';
import { buildReducer, MutationCase } from './reducer.js';
import {
    augmentGetState,
    augmentMiddlewareApi,
    buildEffectsMiddleware,
    EffectMiddleware,
} from './effects.js';
import { ToolkitStore } from '@reduxjs/toolkit/dist/configureStore.js';
import { produce } from 'immer';

type MyActionCreator = { type: string } & ((...args: any) => any);

type XSel<R> = R extends Function ? R : () => R;
type CurriedSelector<S> = S extends (...args: any) => infer R ? XSel<R> : never;

type CurriedSelectors<S> = {
    [key in keyof S]: CurriedSelector<S[key]>;
};

type ResolveAction<
    ActionType extends string,
    ActionArg extends any,
> = ActionArg extends MyActionCreator
    ? ActionArg
    : ActionArg extends (...args: any) => any
    ? ActionCreatorWithPreparedPayload<
          Parameters<ActionArg>,
          ReturnType<ActionArg>,
          ActionType
      >
    : ActionCreatorWithoutPayload<ActionType>;

type ResolveActions<
    A extends {
        [key: string]: any;
    },
> = {
    [ActionType in keyof A]: ActionType extends string
        ? ResolveAction<ActionType, A[ActionType]>
        : never;
};

type Reaction<S = any, M extends MiddlewareAPI = MiddlewareAPI> = (
    api: M,
) => (state: S, previousState: S, unsubscribe: () => void) => any;

type AugmentedMiddlewareAPI<S, A, SELECTORS> = MiddlewareAPI<
    Dispatch<AnyAction>,
    S
> & {
    dispatch: A;
    getState: CurriedSelectors<SELECTORS>;
    actions: A;
    selectors: SELECTORS;
};

export type Mutation<A extends Action<any> = Action<any>, S = any> = (
    payload: A extends {
        payload: infer P;
    }
        ? P
        : undefined,
    action: A,
) => (state: S) => S | void;

type SelectorForState<S> = (state: S) => unknown;
type SelectorsForState<S> = {
    [key: string]: SelectorForState<S>;
};

export default class Updux<
    T_LocalState = Record<any, any>,
    T_LocalActions extends {
        [actionType: string]: any;
    } = {},
    T_Subduxes extends Record<string, Dux> = {},
    T_LocalSelectors extends SelectorsForState<
        AggregateState<T_LocalState, T_Subduxes>
    > = {},
> {
    #localInitial: T_LocalState;
    #localActions: T_LocalActions;
    #localMutations: MutationCase[] = [];
    #defaultMutation: Omit<MutationCase, 'matcher'>;
    #subduxes: T_Subduxes;

    #name: string;

    #actions: AggregateActions<ResolveActions<T_LocalActions>, T_Subduxes>;

    #initialState: AggregateState<T_LocalState, T_Subduxes>;

    #localSelectors: Record<
        string,
        (state: AggregateState<T_LocalState, T_Subduxes>) => any
    >;
    #selectors: any;

    #localEffects: Middleware[] = [];

    constructor(
        config: Partial<{
            initialState: T_LocalState;
            actions: T_LocalActions;
            subduxes: T_Subduxes;
            selectors: T_LocalSelectors;
        }>,
    ) {
        // TODO check that we can't alter the initialState after the fact
        this.#localInitial = config.initialState ?? ({} as T_LocalState);
        this.#localActions = config.actions ?? ({} as T_LocalActions);
        this.#subduxes = config.subduxes ?? ({} as T_Subduxes);

        this.#actions = buildActions(this.#localActions, this.#subduxes);

        this.#initialState = buildInitial(this.#localInitial, this.#subduxes);
        this.#localSelectors = config.selectors;

        const basedSelectors = R.mergeAll(
            Object.entries(this.#subduxes)
                .filter(([slice, { selectors }]) => selectors)
                .map(([slice, { selectors }]) =>
                    R.mapValues(selectors, (s) => (state = {}) => {
                        return s(state?.[slice]);
                    }),
                ),
        );

        this.#selectors = R.merge(basedSelectors, this.#localSelectors);
    }

    get actions() {
        return this.#actions;
    }

    // TODO memoize?
    get initialState() {
        return this.#initialState;
    }

    get effects() {
        return [
            ...this.#localEffects,
            ...Object.entries(this.#subduxes).flatMap(
                ([slice, { effects }]) => {
                    if (!effects) return [];
                    return effects.map(
                        (effect) => (api) =>
                            effect({
                                ...api,
                                getState: () => api.getState()[slice],
                            }),
                    );
                },
            ),
        ];
    }

    get reactions(): any {
        return [
            ...this.#localReactions,
            ...Object.entries(this.#subduxes).flatMap(
                ([slice, { reactions }]) =>
                    reactions.map(
                        (r) => (api, unsub) =>
                            r(
                                {
                                    ...api,
                                    getState: () => api.getState()[slice],
                                },
                                unsub,
                            ),
                    ),
            ),
        ];
    }

    createStore(
        options: Partial<{
            preloadedState: T_LocalState;
        }> = {},
    ) {
        const preloadedState: any = options.preloadedState;

        const effects = buildEffectsMiddleware(
            this.effects,
            this.actions,
            this.selectors,
        );

        const store = configureStore({
            reducer: this.reducer as Reducer<
                AggregateState<T_LocalState, T_Subduxes>,
                AnyAction
            >,
            preloadedState,
            middleware: [effects],
        });

        const dispatch: any = store.dispatch;
        for (const a in this.actions) {
            dispatch[a] = (...args) => {
                const action = (this.actions as any)[a](...args);
                dispatch(action);
                return action;
            };
        }

        store.getState = augmentGetState(store.getState, this.selectors);

        for (const reaction of this.reactions) {
            let unsub;
            const r = reaction(store);

            unsub = store.subscribe(() => r(unsub));
        }

        (store as any).actions = this.actions;
        (store as any).selectors = this.selectors;

        return store as ToolkitStore<AggregateState<T_LocalState, T_Subduxes>> &
            AugmentedMiddlewareAPI<
                AggregateState<T_LocalState, T_Subduxes>,
                AggregateActions<ResolveActions<T_LocalActions>, T_Subduxes>,
                AggregateSelectors<
                    T_LocalSelectors,
                    T_Subduxes,
                    AggregateState<T_LocalState, T_Subduxes>
                >
            >;
    }

    get selectors(): AggregateSelectors<
        T_LocalSelectors,
        T_Subduxes,
        AggregateState<T_LocalState, T_Subduxes>
    > {
        return this.#selectors as any;
    }

    // TODO memoize this sucker
    get reducer() {
        return buildReducer(
            this.initialState,
            this.#localMutations,
            this.#defaultMutation,
            this.#subduxes,
        ) as any as (
            state: undefined | typeof this.initialState,
            action: Action,
        ) => typeof this.initialState;
    }

    // TODO be smarter with the guard?
    addMutation<A extends Action<any>>(
        matcher: (action: A) => boolean,
        mutation: Mutation<A, AggregateState<T_LocalState, T_Subduxes>>,
        terminal?: boolean,
    );
    addMutation<A extends ActionCreator<any>>(
        actionCreator: A,
        mutation: Mutation<
            ReturnType<A>,
            AggregateState<T_LocalState, T_Subduxes>
        >,
        terminal?: boolean,
    );
    addMutation(matcher, mutation, terminal = false) {
        if (typeof matcher === 'function' && matcher.match) {
            // matcher, matcher man...
            matcher = matcher.match;
        }

        const immerMutation = (...args) => produce(mutation(...args));

        this.#localMutations.push({
            terminal,
            matcher,
            mutation: immerMutation,
        });
    }

    addDefaultMutation(
        mutation: Mutation<
            Action<any>,
            AggregateState<T_LocalState, T_Subduxes>
        >,
        terminal = false,
    ) {
        this.#defaultMutation = { mutation, terminal };
    }

    addEffect(
        actionCreator: AggregateActions<
            ResolveActions<T_LocalActions>,
            T_Subduxes
        >,
        effect: EffectMiddleware,
    ): EffectMiddleware;
    addEffect(
        guardFunc: (
            action: AggregateActions<
                ResolveActions<T_LocalActions>,
                T_Subduxes
            >[keyof AggregateActions<
                ResolveActions<T_LocalActions>,
                T_Subduxes
            >],
        ) => boolean,
        effect: EffectMiddleware,
    ): EffectMiddleware;
    addEffect(effect: EffectMiddleware): EffectMiddleware;
    addEffect(...args) {
        let effect;
        if (args.length === 1) {
            effect = args[0];
        } else {
            const [actionCreator, originalEffect] = args;

            const test = actionCreator.hasOwnProperty('match')
                ? actionCreator.match
                : actionCreator;

            effect = (api) => (next) => {
                const e = originalEffect(api)(next);
                return (action) => {
                    const func = test(action) ? e : next;
                    return func(action);
                };
            };
        }

        this.#localEffects.push(effect);

        return effect;
    }

    get effectsMiddleware() {
        return buildEffectsMiddleware(
            this.effects,
            this.actions,
            this.selectors,
        );
    }

    #localReactions: any[] = [];
    addReaction(
        reaction: Reaction<
            AggregateState<T_LocalState, T_Subduxes>,
            AugmentedMiddlewareAPI<
                AggregateState<T_LocalState, T_Subduxes>,
                AggregateActions<ResolveActions<T_LocalActions>, T_Subduxes>,
                AggregateSelectors<
                    T_LocalSelectors,
                    T_Subduxes,
                    AggregateState<T_LocalState, T_Subduxes>
                >
            >
        >,
    ) {
        let previous: any;

        const memoized = (api: any) => {
            api = augmentMiddlewareApi(api, this.actions, this.selectors);
            const r = reaction(api);
            return (unsub: () => void) => {
                const state = api.getState();
                if (state === previous) return;
                let p = previous;
                previous = state;
                r(state, p, unsub);
            };
        };
        this.#localReactions.push(memoized);
    }

    // internal method REMOVE
    subscribeTo(store, subscription) {
        const localStore = augmentMiddlewareApi(
            {
                ...store,
                subscribe: (subscriber) => this.subscribeTo(store, subscriber), // TODO not sure
            },
            this.actions,
            this.selectors,
        );

        const subscriber = subscription(localStore);

        let previous;
        let unsub;

        const memoSub = () => {
            const state = store.getState();
            if (state === previous) return;
            let p = previous;
            previous = state;
            subscriber(state, p, unsub);
        };

        return store.subscribe(memoSub);
    }
}