import React, {
  useMemo,
  useContext,
  useLayoutEffect,
  useState,
  createContext,
} from "react";

export const withUpstreamOverride = <P,>(Component: React.ComponentType<P>) => {
  type ContextValue = [boolean, React.Dispatch<React.SetStateAction<boolean>>];

  const DefaultComponentContext = createContext<ContextValue>([
    false,
    () => {},
  ]);

  const ComponentContext: React.FC<{ children: React.ReactNode }> = ({
    children,
  }) => {
    const [isRenderedUpstream, setIsRenderedUpstream] = useState(false);
    const contextValue: ContextValue = useMemo(
      () => [isRenderedUpstream, setIsRenderedUpstream],
      [isRenderedUpstream],
    );

    return (
      <DefaultComponentContext.Provider value={contextValue}>
        {children}
      </DefaultComponentContext.Provider>
    );
  };

  const DefaultComponent = (
    props: P & {
      // indicates whether component should render when not rendered upstream
      /** @private internal */
      __isFallback?: boolean;
    },
  ) => {
    const [isRenderedUpstream, setIsRenderedUpstream] = useContext(
      DefaultComponentContext,
    );

    useLayoutEffect(() => {
      if (!props.__isFallback) {
        setIsRenderedUpstream(true);
        return () => setIsRenderedUpstream(false);
      }
    }, [props.__isFallback, setIsRenderedUpstream]);

    if (props.__isFallback && isRenderedUpstream) {
      return null;
    }

    return <Component {...props} />;
  };
  if (Component.name) {
    DefaultComponent.displayName = `${Component.name}_upstreamOverrideWrapper`;
    ComponentContext.displayName = `${Component.name}_upstreamOverrideContextWrapper`;
  }

  return [ComponentContext, DefaultComponent] as const;
};