useIonizedBunja.tsx

import type { Bunja, ScopeValuePair } from 'bunja';
import { useBunja } from 'bunja/react';
import type { Atom, createStore, PrimitiveAtom, WritableAtom } from 'jotai';
import { getDefaultStore } from 'jotai';
import { useEffect, useMemo, useReducer } from 'react';

type JotaiStore = ReturnType<typeof createStore>;

export function useIonizedBunja<T extends object>(
  bunja: Bunja<T>,
  scopeValuePairs?: ScopeValuePair<any>[],
  store?: JotaiStore,
) {
  const bunjaInstance = useBunja(bunja, scopeValuePairs);
  const [, rerender] = useReducer((x) => !x, false);
  return useIonizedValue(bunjaInstance, store ?? getDefaultStore(), rerender);
}

export interface ReadonlyIon<T> {
  value: T;
}
export interface WritableIon<T> {
  value: T;
}
export type Ionized<T> = T & {
  [K in keyof T as K extends `${infer U}Atom` ? `${U}Ion` : never]: T[K] extends WritableAtom<
    infer V,
    [any],
    any
  >
    ? WritableIon<V>
    : T[K] extends Atom<infer V>
      ? ReadonlyIon<V>
      : never;
};
function useIonizedValue<T extends object>(value: T, store: JotaiStore, rerender: () => void): Ionized<T> {
  const [unmountCallbacks, ionizedValue] = useMemo(() => {
    const unmountCallbacks: Array<() => void> = [];
    const result = {
      ...value,
      ...Object.fromEntries(
        Object.entries(value).flatMap(([k, maybeAtom]) => {
          if (!k.endsWith('Atom')) return [];

          unmountCallbacks.push(store.sub(maybeAtom, rerender));
          return [
            [
              `${k.slice(0, -4)}Ion`,
              {
                get value() {
                  return store.get(maybeAtom as Atom<any>);
                },
                set value(value: T) {
                  store.set(maybeAtom as PrimitiveAtom<any>, value);
                },
              },
            ],
          ];
        }),
      ),
    } as Ionized<T>;
    return [unmountCallbacks, result];
  }, [rerender, store, value]);
  useEffect(() => {
    return () => {
      unmountCallbacks.forEach((callback) => callback());
    };
  }, [unmountCallbacks]);
  return ionizedValue;
}