import { clsx } from '@trail-ui/shared-utils';
import type { CheckboxSlots, CheckboxVariantProps, SlotsToClasses } from '@trail-ui/theme';
import { checkbox } from '@trail-ui/theme';
import {
  ForwardedRef,
  ReactElement,
  ReactNode,
  cloneElement,
  forwardRef,
  useContext,
  useMemo,
} from 'react';
import type { CheckboxProps as AriaCheckboxProps } from 'react-aria-components';
import { Checkbox as AriaCheckbox } from 'react-aria-components';
import { CheckboxGroupThemeContext } from './checkbox-group';
import { CheckboxIcon, CheckboxIconProps } from './checkbox-icon';

export interface CheckboxProps extends AriaCheckboxProps, CheckboxVariantProps {
  /**
   * The icon to be displayed when the checkbox is checked.
   */
  icon?: ReactNode | ((props: CheckboxIconProps) => ReactNode);
  /**
   * Classes object to style the checkbox and its children.
   */
  classNames?: SlotsToClasses<CheckboxSlots>;
  className?: string;
}

function Checkbox(props: CheckboxProps, ref: ForwardedRef<HTMLLabelElement>) {
  const groupThemeContext = useContext(CheckboxGroupThemeContext);

  const {
    icon = <CheckboxIcon />,
    lineThrough = groupThemeContext?.lineThrough ?? false,
    isInvalid = groupThemeContext?.isInvalid ?? false,
    className,
    classNames,
    children,
    ...otherProps
  } = props;

  const slots = useMemo(
    () =>
      checkbox({
        lineThrough,
        isInvalid,
      }),
    [isInvalid, lineThrough],
  );

  const baseStyles = clsx(classNames?.base, className);

  return (
    <AriaCheckbox ref={ref} className={slots.base({ class: baseStyles })} {...otherProps}>
      {(renderProps) => {
        const iconProps = {
          'data-checked': String(renderProps.isSelected),
          isSelected: renderProps.isSelected,
          isIndeterminate: renderProps.isIndeterminate,
          isInvalid: renderProps.isInvalid,
          className: slots.icon({ class: classNames?.icon }),
        };
        const clonedIcon =
          typeof icon === 'function'
            ? icon(iconProps)
            : cloneElement(icon as ReactElement, iconProps);

        return (
          <>
            <span
              aria-hidden="true"
              className={slots.control({ class: classNames?.control })}
              data-indeterminate={props.isIndeterminate}
            >
              {clonedIcon}
            </span>
            {children && (
              <span className={slots.label({ class: classNames?.label })}>
                {typeof children === 'function' ? children(renderProps) : children}
              </span>
            )}
          </>
        );
      }}
    </AriaCheckbox>
  );
}

/**
 * A checkbox allows a user to select multiple items from a list of individual items, or
 * to mark one individual item as selected.
 */
const _Checkbox = forwardRef(Checkbox);

export { _Checkbox as Checkbox };