import { User, UserAccount } from "@frontend/kitui";
import { produce } from "immer";
import { create } from "zustand";
import { createJSONStorage, devtools, persist } from "zustand/middleware";

interface UserStore {
  userId: string | null;
  user: User | null;
  userAccount: UserAccount | null;
}

const initialState: UserStore = {
  userId: null,
  user: null,
  userAccount: null,
};

export const useUserStore = create<UserStore>()(
  persist(
    devtools((set, get) => initialState, {
      name: "User",
      enabled: process.env.NODE_ENV === "development",
      trace: process.env.NODE_ENV === "development",
    }),
    {
      version: 2,
      name: "user",
      partialize: (state) => ({
        userId: state.userId,
        user: state.user,
      }),
      migrate: (persistedState, version) => ({
        ...(persistedState as UserStore),
        user: null,
      }),
    },
  ),
);

export const setUserId = (userId: string | null) =>
  useUserStore.setState({ userId });
export const setUser = (user: User) =>
  useUserStore.setState(
    produce<UserStore>((state) => {
      state.user = user;
    }),
  );

export const clearUserData = () => useUserStore.setState({ ...initialState });

export const setUserAccount = (userAccount: UserAccount) =>
  useUserStore.setState({ userAccount });