import { useStoreState } from "@ariakit/react-store";
import {
  useEvent,
  useForceUpdate,
  useMergeRefs,
  useTagName,
  useWrapElement,
  createElement,
  createHook,
  forwardRef,
} from "@ariakit/react-utils";
import type { Props } from "@ariakit/react-utils";
import { disabledFromProps, removeUndefinedValues } from "@ariakit/utils";
import type {
  ChangeEvent,
  ComponentPropsWithoutRef,
  ElementType,
  MouseEvent,
} from "react";
import { useEffect, useRef, useState } from "react";
import type { CommandOptions } from "../command/command.tsx";
import { useCommand } from "../command/command.tsx";
import { getPrimitiveValue } from "./__utils.ts";
import { CheckboxCheckedContext } from "./checkbox-checked-context.tsx";
import { useCheckboxContext } from "./checkbox-context.tsx";
import type { CheckboxStore } from "./checkbox-store.ts";

const TagName = "input" satisfies ElementType;
type TagName = typeof TagName;
type HTMLType = HTMLElementTagNameMap[TagName];

function setMixed(element: HTMLType, mixed?: boolean) {
  if (mixed) {
    element.indeterminate = true;
  } else if (element.indeterminate) {
    element.indeterminate = false;
  }
}

function isNativeCheckbox(tagName?: string, type?: string) {
  return tagName === "input" && (!type || type === "checkbox");
}

/**
 * Returns props to create a `Checkbox` component. If the element is not a
 * native checkbox, the hook will return additional props to make sure it's
 * accessible.
 * @see https://ariakit.com/components/checkbox
 * @example
 * ```jsx
 * const props = useCheckbox({ render: <div /> });
 * <Role {...props}>Accessible checkbox</Role>
 * ```
 */
export const useCheckbox = createHook<TagName, CheckboxOptions>(
  function useCheckbox({
    store,
    name,
    value: valueProp,
    checked: checkedProp,
    defaultChecked,
    ...props
  }) {
    const context = useCheckboxContext();
    store = store || context;

    const [_checked, setChecked] = useState(defaultChecked ?? false);

    const checked = useStoreState(store, ["value"], (state) => {
      if (checkedProp !== undefined) return checkedProp;
      if (state?.value === undefined) return _checked;
      if (valueProp != null) {
        if (Array.isArray(state.value)) {
          const primitiveValue = getPrimitiveValue(valueProp);
          return state.value.includes(primitiveValue);
        }
        return state.value === valueProp;
      }
      if (Array.isArray(state.value)) return false;
      if (typeof state.value === "boolean") return state.value;
      return false;
    });

    const ref = useRef<HTMLType>(null);
    const tagName = useTagName(ref, TagName);
    const nativeCheckbox = isNativeCheckbox(tagName, props.type);
    const mixed = checked ? checked === "mixed" : undefined;
    const isChecked = checked === "mixed" ? false : checked;
    const disabled = disabledFromProps(props);
    // When the checked property is programmatically set on the change event, we
    // need to schedule the element's property update, so the controlled
    // isChecked state can be taken into account.
    const [propertyUpdated, schedulePropertyUpdate] = useForceUpdate();

    useEffect(() => {
      const element = ref.current;
      if (!element) return;
      setMixed(element, mixed);
      if (nativeCheckbox) return;
      element.checked = isChecked;
      if (name !== undefined) {
        element.name = name;
      }
      if (valueProp !== undefined) {
        element.value = String(valueProp);
      }
    }, [propertyUpdated, mixed, nativeCheckbox, isChecked, name, valueProp]);

    const onChangeProp = props.onChange;

    const onChange = useEvent((event: ChangeEvent<HTMLType>) => {
      if (disabled) {
        event.stopPropagation();
        event.preventDefault();
        return;
      }
      setMixed(event.currentTarget, mixed);
      if (!nativeCheckbox) {
        // If the element is not a native checkbox, we need to manually update
        // its checked property.
        event.currentTarget.checked = !event.currentTarget.checked;
        schedulePropertyUpdate();
      }
      onChangeProp?.(event);
      if (event.defaultPrevented) return;

      const elementChecked = event.currentTarget.checked;
      setChecked(elementChecked);

      store?.setValue((prevValue) => {
        if (valueProp == null) return elementChecked;
        const primitiveValue = getPrimitiveValue(valueProp);
        if (!Array.isArray(prevValue)) {
          return prevValue === primitiveValue ? false : primitiveValue;
        }
        if (elementChecked) {
          if (prevValue.includes(primitiveValue)) {
            return prevValue;
          }
          return [...prevValue, primitiveValue];
        }
        return prevValue.filter((v) => v !== primitiveValue);
      });
    });

    const onClickProp = props.onClick;

    const onClick = useEvent((event: MouseEvent<HTMLType>) => {
      onClickProp?.(event);
      if (event.defaultPrevented) return;
      if (nativeCheckbox) return;
      // @ts-expect-error The onChange event expects a ChangeEvent, but here we
      // need to pass a MouseEvent.
      onChange(event);
    });

    props = useWrapElement(
      props,
      (element) => (
        <CheckboxCheckedContext.Provider value={isChecked}>
          {element}
        </CheckboxCheckedContext.Provider>
      ),
      [isChecked],
    );

    props = {
      role: !nativeCheckbox ? "checkbox" : undefined,
      type: nativeCheckbox ? "checkbox" : undefined,
      "aria-checked": checked,
      ...props,
      ref: useMergeRefs(ref, props.ref),
      onChange,
      onClick,
    };

    props = useCommand<TagName>({ clickOnEnter: !nativeCheckbox, ...props });

    // useFormCheckbox computes the field name after this hook. Remove the
    // native-only sentinels before they can overwrite that later default.
    return removeUndefinedValues({
      name: nativeCheckbox ? name : undefined,
      value: nativeCheckbox ? valueProp : undefined,
      checked: isChecked,
      ...props,
    });
  },
);

/**
 * Renders an accessible checkbox element. If the underlying element is not a
 * native checkbox, this component will pass additional attributes to make sure
 * it's accessible.
 * @see https://ariakit.com/components/checkbox
 * @example
 * ```jsx
 * <Checkbox render={<div />}>Accessible checkbox</Checkbox>
 * ```
 */
export const Checkbox = forwardRef(function Checkbox(props: CheckboxProps) {
  const htmlProps = useCheckbox(props);
  return createElement(TagName, htmlProps);
});

export interface CheckboxOptions<
  T extends ElementType = TagName,
> extends CommandOptions<T> {
  /**
   * Object returned by the
   * [`useCheckboxStore`](https://ariakit.com/reference/use-checkbox-store)
   * hook. If not provided, the closest
   * [`CheckboxProvider`](https://ariakit.com/reference/checkbox-provider)
   * component's context will be used. Otherwise, the component will fall back
   * to an internal store.
   *
   * Live examples:
   * - [Checkbox as button](https://ariakit.com/examples/checkbox-as-button)
   */
  store?: CheckboxStore;
  /**
   * The native `name` attribute.
   *
   * Live examples:
   * - [MenuItemCheckbox](https://ariakit.com/examples/menu-item-checkbox)
   */
  name?: string;
  /**
   * The value of the checkbox. This is useful when the same checkbox store is
   * used for multiple [`Checkbox`](https://ariakit.com/reference/checkbox)
   * elements, in which case the value will be an array of checked values.
   *
   * Live examples:
   * - [Checkbox group](https://ariakit.com/examples/checkbox-group)
   * - [MenuItemCheckbox](https://ariakit.com/examples/menu-item-checkbox)
   * @example
   * ```jsx "value"
   * <CheckboxProvider defaultValue={["Apple", "Orange"]}>
   *   <Checkbox value="Apple" />
   *   <Checkbox value="Orange" />
   *   <Checkbox value="Watermelon" />
   * </CheckboxProvider>
   * ```
   */
  value?: ComponentPropsWithoutRef<TagName>["value"];
  /**
   * The default checked state of the checkbox. This prop is ignored if the
   * [`checked`](https://ariakit.com/reference/checkbox#checked) or the
   * [`store`](https://ariakit.com/reference/checkbox#store) props are provided.
   */
  defaultChecked?: "mixed" | boolean;
  /**
   * The checked state of the checkbox. This will override the value inferred
   * from [`store`](https://ariakit.com/reference/checkbox#store) prop, if
   * provided. This can be `"mixed"` to indicate that the checkbox is partially
   * checked.
   */
  checked?: "mixed" | boolean;
  /**
   * A function that is called when the checkbox's checked state changes.
   */
  onChange?: ComponentPropsWithoutRef<TagName>["onChange"];
}

export type CheckboxProps<T extends ElementType = TagName> = Props<
  T,
  CheckboxOptions<T>
>;