import React, { createContext, useCallback, useContext, type FC, type ReactNode } from 'react';
import type { ArrayPath, Path, PathValue } from 'react-hook-form';
import { get, set } from '@wener/utils';
import { create as produce } from 'mutative';
import { createStore, useStore, type StoreApi } from 'zustand';

const DefaultContextStore = createStore(() => {
  return {};
});

const Context = createContext<StoreApi<any> | undefined>(undefined);

export const ContextStoreProvider: FC<{ value: StoreApi<any>; children?: ReactNode }> = ({ value, children }) => {
  return <Context.Provider value={value}>{children}</Context.Provider>;
};

function getContextStore() {
  return DefaultContextStore;
}

export function useContextStore<O extends Record<string, any>>(): UseContextStoreReturn<O> {
  const store = useContext(Context) ?? DefaultContextStore;
  return {
    store,
    set(path: string, value: any) {
      store.setState(
        produce((s: any) => {
          set(s, path, value, false);
        }),
      );
    },
    get(path: string): any {
      return get(store.getState(), path);
    },
    useWatch(path: string): any {
      return useStore(
        store,
        useCallback((s) => get(s, path), [path]),
      );
    },
  };
}

export interface UseContextStoreReturn<O extends Record<string, any> = Record<string, any>> {
  store: StoreApi<any>;

  set<P extends Path<O> | ArrayPath<O>, V extends PathValue<O, P>>(type: P, payload: V): void;

  get<P extends Path<O> | ArrayPath<O>, V extends PathValue<O, P>>(path: P): V;

  useWatch<P extends Path<O> | ArrayPath<O>, V extends PathValue<O, P>>(path: P): V;
}