// deno-lint-ignore-file no-explicit-any

import type { StateStorage } from "zustand/middleware";

type MaybeAsync<U> = U | Promise<U>;

export type SerializerConfig<TValue> = {
  serializer: (value: TValue) => any | Promise<any>;
  deserializer: (value: unknown) => TValue | Promise<TValue>;
  initializer?: () => MaybeAsync<TValue>;
};

export type SerializerMap<T> = {
  [K in keyof T]?: SerializerConfig<T[K]>;
};

export function makeSlicedSerializerStorage<T extends object>(
  storage: StateStorage,
  serializers: SerializerMap<T>,
): StateStorage {
  // serialize by looping through the state and using registered serializers per key
  async function serializeState(state: T) {
    const output = {} as Record<string, unknown>;
    for (const key in state) {
      let value = await state[key];
      if (serializers[key]) {
        value = await serializers[key].serializer(value);
      }

      output[key] = value;
    }

    return output;
  }

  // deserialize by looping through the state and using registered deserializers per key
  async function deserializeState(state: any) {
    const output = {} as T;
    for (const key_ in state) {
      const key = key_ as keyof T;
      output[key] = state[key];
      if (serializers[key]) {
        output[key] = await serializers[key].deserializer(output[key]);
      }
    }

    return output;
  }

  // initialize also run here, so we can provide default initial content as if it's coming from the store
  async function initializeState() {
    const state = {} as T;
    for (const key in serializers) {
      const initializer = serializers[key]?.initializer;
      if (initializer) {
        state[key] = await initializer();
      }
    }

    return state as T;
  }

  return {
    async getItem(name) {
      const result = await storage.getItem(name);
      if (result == null) {
        return { version: 0, state: await initializeState() };
      }

      const { version, state } = result;
      return { version, state: await deserializeState(state) };
    },

    async setItem(name, input) {
      const { version, state } = input;

      storage.setItem(name, { version, state: await serializeState(state) });
    },

    removeItem(name) {
      storage.removeItem(name);
    },
  };
}