import { create, type StoreApi, type UseBoundStore } from "zustand";
import { createJSONStorage, devtools, persist } from "zustand/middleware";

import {
  type AudioSlice,
  createAudioSlice,
  createDebugSlice,
  createLayoutSlice,
  type DebugSlice,
  type LayoutSlice,
} from "./slices";

// Auto generate selectors
// https://docs.pmnd.rs/zustand/guides/auto-generating-selectors
type WithSelectors<S> = S extends { getState: () => infer T }
  ? S & { use: { [K in keyof T]: () => T[K] } }
  : never;

const createSelectors = <S extends UseBoundStore<StoreApi<object>>>(
  _store: S,
) => {
  const store = _store as WithSelectors<typeof _store>;
  store.use = {};
  for (const k of Object.keys(store.getState())) {
    (store.use as any)[k] = () => store((s) => s[k as keyof typeof s]);
  }

  return store;
};

export interface CombinedState {
  audio: AudioSlice;
  debug: DebugSlice;
  layout: LayoutSlice;
}

type AppState = AudioSlice & DebugSlice & LayoutSlice;

// Why the currying ()(...)?
// TLDR: It is a workaround for microsoft/TypeScript#10571.
const useBoundStoreBase = create<AppState>()(
  devtools(
    persist(
      (...a) => ({
        ...createAudioSlice(...a),
        ...createDebugSlice(...a),
        ...createLayoutSlice(...a),
      }),
      {
        name: "persisted-store",
        storage: createJSONStorage(() => sessionStorage),
        partialize: (state) => ({
          isMuted: state.isMuted,
        }),
      },
    ),
  ),
);

export const useStore = createSelectors(useBoundStoreBase);