import * as Ariakit from '@ariakit/react';
import { ComponentPropsWithRef, forwardRef, useContext } from 'react';
import styled from 'styled-components';
import { UIComponentWithRef } from '../../components/types';
import { getVariantKey, useVariants } from '../../variants';
import { CheckboxContext } from './checkbox.context';
import { checkboxVariants, ICheckboxVariants } from './checkbox.variants';

type Props = {
  variant?: keyof typeof checkboxVariants | (string & {});
  variants?: ICheckboxVariants;
};

// export type CheckboxGroupProps = UIStyledComponentProps<Props, 'input'>;

const StyledCheckboxGroup = styled(Ariakit.Group) <Ariakit.GroupProps & Props>`
  ${props => getVariantKey('group', props.variant, props.variants)}
`;

type CheckboxGroupComponent<C extends React.ElementType = 'div'> = UIComponentWithRef<C, Ariakit.CheckboxProviderProps & Props>;

export const CheckboxGroup: CheckboxGroupComponent = forwardRef<HTMLDivElement, CheckboxGroupProps>(
  function CheckboxGroup({
    variant = 'default',
    variants,
    defaultValue,
    value,
    setValue,
    store,
    children,
    ...props
  }, ref) {
    const context = useContext(CheckboxContext);
    const globalVariants = useVariants('checkbox', checkboxVariants);
    variants = variants || globalVariants;
    return (
      <Ariakit.CheckboxProvider
        defaultValue={defaultValue}
        value={value}
        setValue={setValue}
        store={store}
      >
        <CheckboxContext.Provider value={{
          variant,
          variants,
        }}>
          <StyledCheckboxGroup
            ref={ref}
            variant={variant || context?.variant}
            variants={variants || context?.variants}
            {...props}
          >
            {children}
          </StyledCheckboxGroup>
        </CheckboxContext.Provider>
      </Ariakit.CheckboxProvider>
    );
  }
);

export type CheckboxGroupProps = ComponentPropsWithRef<typeof CheckboxGroup>;