import React, { createContext, useContext, ReactNode, useMemo } from "react";
import { create, UseBoundStore } from "zustand";

interface ZStoreContextType {
  stores: Map<string, UseBoundStore<any>>;
  registerStore: (name: string, store: UseBoundStore<any>) => void;
  getStore: (name: string) => UseBoundStore<any> | undefined;
}

const ZStoreContext = createContext<ZStoreContextType | null>(null);

interface ZStoreProviderProps {
  children: ReactNode;
  name: string;
  value: any;
}

export const ZStoreProvider: React.FC<ZStoreProviderProps> = ({
  children,
  name,
  value,
}) => {
  const parentContext = useContext(ZStoreContext);

  const context = useMemo(() => {
    if (parentContext) {
      // If parent context exists, create a new store and register it
      const store = create(() => value);
      parentContext.registerStore(name, store);

      return parentContext;
    }

    // Create a new context if none exists
    const stores = new Map<string, UseBoundStore<any>>();

    // Create a store from the provided value
    const store = create(() => value);
    stores.set(name, store);

    return {
      stores,
      registerStore: (storeName: string, storeInstance: UseBoundStore<any>) => {
        stores.set(storeName, storeInstance);
      },
      getStore: (storeName: string) => stores.get(storeName),
    };
  }, [parentContext, name, value]);

  return (
    <ZStoreContext.Provider value={context}>{children}</ZStoreContext.Provider>
  );
};

export const useZStore = (name: string): UseBoundStore<any> => {
  const context = useContext(ZStoreContext);

  if (!context) {
    throw new Error("useZStore must be used within a ZStoreProvider");
  }

  const store = context.getStore(name);

  if (!store) {
    throw new Error(
      `Store with name "${name}" not found. Make sure it's provided by a ZStoreProvider.`
    );
  }

  return store;
};