"use client"
import React, {
  Children,
  ComponentProps,
  ElementRef,
  forwardRef,
  useLayoutEffect,
  useMemo,
  useState,
} from "react"
import { createContext, useContextSelector } from "use-context-selector"
import { ChevronLeft } from "../../icons"
import { InputVariantProps, inputVariants } from "../../styles"
import { classNames } from "../../utils"
import {
  Dropdown,
  DropdownItem,
  DropdownItemAvatar,
  DropdownItemContent,
  DropdownItemDescription,
  DropdownItemTitle,
  DropdownProps,
  DropdownSize,
} from "../Dropdown"
import { UnstyledButton, UnstyledButtonProps } from "../UnstyledButton"

type SelectPropsBase = Omit<DropdownProps, "content" | "variant" | "size"> &
  InputVariantProps & {
    name?: string
    placeholder?: React.ReactNode
    onValueChange?: (value: string) => void
    overrides?: {
      Content?: {
        className?: string
      }
    }
  }

type SelectPropsSingleSelection = SelectPropsBase & {
  selection?: "single"
  /**
   * The value of the Select. Value must be a subset of the value of the SelectItems rendered as children.
   * In multiple selection mode, this should be undefined.
   */
  value?: string
  /**
   * A function that renders the current value of the Select. Defaults to value's SelectItem.Title or the value itself.
   * In multiple selection mode, this should be undefined.
   */
  renderValue?: (value: string) => React.ReactNode
}

type SelectPropsMultipleSelection = SelectPropsBase & {
  /**
   * When selection mode is multiple, the Select does not close when an item is selected.
   */
  selection: "multiple"
  /**
   * Placeholder is required when selection mode is multiple.
   */
  placeholder: string
  value?: undefined
  renderValue?: undefined
}

type SelectSize = "xs" | "sm" | "md" | "lg"

export type SelectProps = (
  | SelectPropsSingleSelection
  | SelectPropsMultipleSelection
) & { size?: SelectSize; dropdownSize?: DropdownSize }

type SelectContextType = Pick<SelectProps, "selection" | "onValueChange"> & {
  itemRefs: HTMLButtonElement[]
  setItemRefs: React.Dispatch<React.SetStateAction<HTMLButtonElement[]>>
  forceMount: true | undefined
  setForceMount: (forceMount: true | undefined) => void
}

const SelectContext = createContext<SelectContextType>({
  selection: "single",
  itemRefs: [],
  setItemRefs: () => {},
  onValueChange: () => {},
  forceMount: undefined,
  setForceMount: () => {},
})

const SelectBase = forwardRef<
  ElementRef<typeof Dropdown>,
  Omit<SelectProps, "onValueChange" | "selection">
>(function Select(
  {
    value,
    children,
    className,
    placeholder,
    disabled,
    error,
    id,
    name,
    renderValue,
    variant = "default",
    size = "sm",
    dropdownSize = "sm",
    ...props
  },
  ref,
) {
  const [itemRefs, setItemRefs, forceMount, setForceMount] = useContextSelector(
    SelectContext,
    v => [v.itemRefs, v.setItemRefs, v.forceMount, v.setForceMount],
  )

  // This is a hack to force the dropdown to initially mount on render so we can properly
  // set the itemRefs. This is so we can render the correct value in the trigger.
  // Furthermore, we re-mount when children change/values are added.
  useLayoutEffect(() => {
    setForceMount(undefined)
  }, [forceMount, setForceMount])

  useLayoutEffect(() => {
    // Clean up item refs every time children change so we don't accumulate refs indefinitely
    setItemRefs([])

    if (forceMount) {
      setForceMount(undefined)
    } else {
      setForceMount(true)
    }
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [Children.toArray(children).length])

  const displayText = useMemo(() => {
    return itemRefs
      .find(el => el.value === value)
      ?.querySelector("[data-item-text]")?.textContent
  }, [itemRefs, value])

  return (
    <Dropdown
      content={children}
      forceMount={forceMount}
      key={forceMount ? "force-mount" : "normal"}
      ref={ref}
      size={dropdownSize}
      {...props}
    >
      <SelectButton
        className={className}
        disabled={disabled}
        error={error}
        id={id}
        name={name}
        size={size}
        variant={variant}
      >
        {renderValue && value
          ? renderValue(value)
          : displayText || value || placeholder}
      </SelectButton>
    </Dropdown>
  )
})

const SelectBaseWithContext = forwardRef<
  ElementRef<typeof SelectBase>,
  SelectProps
>(function SelectWithContext(
  { selection = "single", onValueChange, ...rest },
  ref,
) {
  const [itemRefs, setItemRefs] = useState<HTMLButtonElement[]>([])
  const [forceMount, setForceMount] = useState<true | undefined>(true)

  return (
    <SelectContext.Provider
      value={{
        selection,
        itemRefs,
        setItemRefs,
        onValueChange,
        forceMount,
        setForceMount,
      }}
    >
      <SelectBase {...rest} ref={ref} />
    </SelectContext.Provider>
  )
})

type SelectItemProps = ComponentProps<typeof DropdownItem> & {
  value: string
}

export type SelectButtonProps = InputVariantProps & UnstyledButtonProps

export const SelectButton = forwardRef<
  ElementRef<typeof UnstyledButton>,
  SelectButtonProps
>(function SelectButton(
  { children, className, disabled, error, variant, size, ...rest },
  ref,
) {
  return (
    <UnstyledButton
      className={classNames(
        inputVariants({ disabled, error, variant, size, rightIcon: true }),
        "justify-between focus:outline-hidden",
        "[&[data-state=closed]>svg]:-rotate-90 [&[data-state=open]>svg]:rotate-90",
        className,
      )}
      disabled={disabled}
      ref={ref}
      {...rest}
    >
      {children}
      <ChevronLeft
        className="duration-[time:250ms] transition"
        size={size === "xs" ? 16 : size === "sm" ? 20 : 24}
      />
    </UnstyledButton>
  )
})

const SelectItemBase = forwardRef<
  ElementRef<typeof DropdownItem>,
  SelectItemProps
>(function SelectItem({ value, onSelect, ...rest }, ref) {
  const [itemRefs, setItemRefs, onValueChange, selection] = useContextSelector(
    SelectContext,
    v => [v.itemRefs, v.setItemRefs, v.onValueChange, v.selection],
  )

  return (
    <DropdownItem
      {...rest}
      ref={(el: HTMLDivElement & HTMLButtonElement) => {
        if (ref) {
          if (typeof ref === "function") {
            ref(el)
          } else {
            ref.current = el
          }
        }

        if (
          el instanceof HTMLButtonElement &&
          !itemRefs.map(item => item.value).includes(el.value)
        ) {
          el.value = value
          setItemRefs(itemRefs => [...itemRefs, el])
        }
      }}
      onSelect={e => {
        onSelect?.(e)
        onValueChange?.(value)
        // Prevent the dropdown from being closed when selecting an item for the multiple case
        if (selection === "multiple") {
          e.preventDefault()
        }
      }}
    />
  )
})

export const SelectItemAvatar = DropdownItemAvatar

export const SelectItemTitle = (
  props: ComponentProps<typeof DropdownItemTitle>,
) => {
  return <DropdownItemTitle data-item-text {...props} />
}

export const SelectItemDescription = DropdownItemDescription

export const SelectItemContent = DropdownItemContent

export const SelectItem = Object.assign(SelectItemBase, {
  Avatar: SelectItemAvatar,
  Title: SelectItemTitle,
  Description: SelectItemDescription,
  Content: SelectItemContent,
})

/**
 * Displays a list of options for the user to pick from — triggered by a button.
 */
export const Select = Object.assign(SelectBaseWithContext, {
  Item: SelectItem,
})