import hoistNonReactStatics from 'hoist-non-react-statics';
import React, { Context, useContext } from 'react';
import { Subtract } from 'utility-types';

// -- TYPES

interface IWithContextOpts<TContext> {
  context: Context<TContext>;
  keys?: Array<keyof TContext>;
}

// -- MAIN

function withContext<
  TProps extends Partial<TContext>,
  TContext extends {},
  TComponent extends React.ComponentType<TProps> = React.ComponentType<TProps>
>({ context, keys }: IWithContextOpts<TContext>) {
  return function enhancer(Component: TComponent) {
    function WithContext(props: Subtract<TProps, Partial<TContext>>) {
      const contextValue = useContext(context);
      let result: Partial<TContext> = contextValue;

      if (keys) {
        result = {};

        for (const key of keys) {
          result[key] = contextValue[key];
        }
      }

      // @ts-ignore TODO: resolve type conflict
      return <Component {...(props as TProps)} {...result} />;
    }

    return hoistNonReactStatics(WithContext, Component);
  };
}

export default withContext;