import * as React from "react";
import type { StoreApi } from "zustand";
import { devtools } from "zustand/middleware";
import {
  UseBoundStoreWithEqualityFn,
  createWithEqualityFn as create,
} from "zustand/traditional";
import { useSuffixContext } from "../contexts";
import { shallow } from "../utils";
import { omitCommonProps } from "../utils/props";

export type State = Record<string, Record<string, unknown>> & {
  reset: (id: string) => void;
  set: (
    id: string,
    newStateForId: Record<string, unknown>,
    name?: string,
  ) => void;
};

export const ZustandContext = React.createContext<
  UseBoundStoreWithEqualityFn<StoreApi<State>>
>({} as any);

export function StoreProvider({ store, children }) {
  return (
    <ZustandContext.Provider value={store}>{children}</ZustandContext.Provider>
  );
}

type StateSelector<State, StateSlice> = (state: State) => StateSlice;
type EqualityChecker<StateSlice> = (a: StateSlice, b: StateSlice) => boolean;

export function useStore<StateSlice = State>(
  selector: StateSelector<State, StateSlice>,
  equalityFn?: EqualityChecker<StateSlice>,
) {
  const useZStore = React.useContext(ZustandContext);
  return useZStore(selector, equalityFn);
}

export const createDebugStore = (
  initialValues: Record<string, Record<string, unknown>> = {},
) =>
  create(
    devtools(
      (set) =>
        ({
          ...initialValues,
          reset: (id) => {
            set(
              (oldState: State) => {
                const { [id]: _, ...rest } = oldState;
                return rest;
              },
              true,
              `reset ${id}`,
            );
          },
          set: (
            id: string,
            newStateForId: Record<string, unknown>,
            name?: string,
          ) => {
            set(
              (oldState: State) => {
                const mergedDataForId = Object.assign(
                  {},
                  oldState[id],
                  typeof newStateForId === "function"
                    ? (newStateForId as any)(oldState[id])
                    : newStateForId,
                );
                return Object.assign({}, oldState, {
                  [id]: mergedDataForId,
                });
              },
              false,
              name,
            );
          },
        }) as State,
      { name: "BrevityStore" },
    ),
    shallow,
  );

export const createStore = (
  initialValues: Record<string, Record<string, unknown>> = {},
) =>
  create(
    (set) =>
      ({
        ...initialValues,
        reset: (id) => {
          set((oldState: State) => {
            const { [id]: _, ...rest } = oldState;
            return rest;
          }, true);
        },
        set: (
          id: string,
          newStateForId: Record<string, unknown>,
          name?: string,
        ) => {
          set((oldState: State) => {
            const mergedDataForId = Object.assign(
              {},
              oldState[id],
              typeof newStateForId === "function"
                ? (newStateForId as any)(oldState[id])
                : newStateForId,
            );
            return Object.assign({}, oldState, {
              [id]: mergedDataForId,
            });
          }, false);
        },
      }) as State,
    shallow,
  );

const DEFAULT_STATE = {};
export const useGet = (
  ids: string[],
  getKey = (id: string) => id,
  suffix?: string,
) => {
  const selector = (state: State) => {
    return ids.reduce(
      (acc, id) => {
        // I'm not sure if we actually need the suffix piece. This is here to preserve old haviour which was erronious I think but I'm scared I'll break an app
        // relying this this behviour so I'm leaving it in.
        acc[id] =
          state[id] ??
          state[getKey(id)] ??
          (suffix ? state[`${id}-${suffix}`] : DEFAULT_STATE) ??
          DEFAULT_STATE;
        return acc;
      },
      {} as Record<string, unknown>,
    );
  };
  return useStore(selector, (a, b) => shallow(a, b, false)) as Record<
    string,
    unknown
  >;
};

export function useGetSet<T = Record<string, unknown>>(
  id: string,
  data: T = DEFAULT_STATE as T,
  disableReset = false,
): [T, (value: Partial<T>, name?: string) => void, () => void] {
  const initialData = React.useRef(data);
  const selector = React.useCallback(
    (state: State) => {
      const setter = state.set;
      const reset = state.reset;
      if (id in state) {
        return [state[id], setter, reset, true];
      }
      const newState = Object.assign({}, data);
      return [newState, setter, reset, false] as const;
    },
    [id, data],
  );
  const [value, set, reset, found] = useStore(selector, shallow) as any;
  const resetLocal = React.useCallback(() => {
    reset(id);
  }, [id, reset]);
  React.useEffect(() => {
    if (!found) {
      set(
        id,
        data,
        process.env.PREVIEW ? `initialize: ${id.slice(-4)}` : undefined,
      );
    }
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [found, id]);
  React.useEffect(() => {
    if (!shallow(data, initialData.current, true)) {
      initialData.current = data;
      set(
        id,
        data,
        process.env.PREVIEW ? `update: ${id.slice(-4)}` : undefined,
      );
    }
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [data]);
  React.useEffect(() => {
    if (disableReset) {
      return;
    }
    return resetLocal;
  }, [resetLocal, disableReset]);
  const setter = React.useCallback(
    (newStateForId: Record<string, unknown>, name?: string) => {
      set(
        id,
        newStateForId,
        process.env.PREVIEW ? `${name} (${id.slice(-4)})` : undefined,
      );
    },
    // eslint-disable-next-line react-hooks/exhaustive-deps
    [id],
  );
  return [value as T, setter, resetLocal];
}

export const StateProvider = ({
  children,
  ids,
  className,
  getKey = (id) => id,
  ...props
}: {
  ids: string[];
  getKey?: (id: string) => string;
  children: (state?: Record<string, unknown>) => React.ReactNode;
  className?: string;
  ref?: React.Ref<any>;
}) => {
  const suffix = useSuffixContext();
  const kids = children(useGet(ids, getKey, suffix));
  return (
    <>
      {React.Children.map(kids, (child) =>
        React.cloneElement(child as any, omitCommonProps(props)),
      )}
    </>
  );
};