blitzkit/packages/varuna/src/index.ts

import { produce } from 'immer';
import { merge } from 'lodash-es';
import { useEffect, useState } from 'react';

interface ProviderProps<Arguments> {
  args: Arguments;
}

export class Varuna<Type, Arguments = void> {
  private listeners = new Set<(state: Type) => void>();

  private initialized = false;
  private _initial?: Type;
  private _state?: Type;

  constructor(
    private creator: Type | ((...args: Arguments[]) => Type),
    private persistence?: string,
  ) {
    if (typeof creator !== 'function') this.initialize(creator);
  }

  Provider({ args }: ProviderProps<Arguments>) {
    if (typeof this.creator !== 'function') {
      throw new Error('Provider must be used with a function creator');
    }

    this.initialize((this.creator as (...args: Arguments[]) => Type)(args));

    return null;
  }

  private initialize(initial: Type) {
    let data = initial;

    if (this.persistence) {
      const dehydrated = localStorage.getItem(this.persistence);

      if (dehydrated) {
        const rehydrated = JSON.parse(dehydrated) as unknown;
        data = merge(data, rehydrated);
      }
    }

    this.initialized = true;
    this._initial = data;
    this._state = data;
  }

  private dispatch() {
    const state = this.state;
    this.listeners.forEach((callback) => callback(state));
  }

  on<Slice>(slicer: (state: Type) => Slice, callback: (state: Slice) => void) {
    const last = slicer(this.state);

    function listener(state: Type) {
      const next = slicer(state);
      if (next !== last) callback(next);
    }

    this.listeners.add(listener);

    return () => {
      this.listeners.delete(listener);
    };
  }

  mutate(mutator: (state: Type) => void) {
    this.set(produce(this.state, mutator));
  }

  private assertInitialized() {
    if (!this.initialized) throw new Error('State used before initialization');
  }

  get state() {
    this.assertInitialized();
    return this._state as Type;
  }

  set(state: Type) {
    this.assertInitialized();
    this._state = state;
    this.dispatch();
  }

  get initial() {
    this.assertInitialized();
    return this._initial as Type;
  }

  use<Slice = Type>(
    slicer: (state: Type) => Slice = (state) => state as unknown as Slice,
  ) {
    const [slice, setSlice] = useState(slicer(this.state));
    useEffect(() => this.on(slicer, setSlice), []);
    return slice;
  }
}