'use client'

import React from 'react'
import { cn } from '../../utils'
import { LabeledCheckbox } from './LabeledCheckbox'

// Context types
interface NestedCheckboxContextValue {
  /* The identifiers of currently selected option checkboxes */
  selectedIds: string[]
  /* All option identifiers that are currently mounted within this NestedCheckbox */
  optionIds: string[]
  /* Register an option when it mounts */
  registerOption: (id: string) => void
  /* Unregister an option when it unmounts */
  unregisterOption: (id: string) => void
  /* Toggle a single option */
  toggleOption: (id: string, checked: boolean) => void
  /* Select all options */
  selectAll: () => void
  /* Deselect all options */
  deselectAll: () => void
  /* Disabled state propagated from provider */
  disabled?: boolean
}

const NestedCheckboxContext =
  React.createContext<NestedCheckboxContextValue | null>(null)

function useNestedCheckboxContext(componentName: string) {
  const ctx = React.useContext(NestedCheckboxContext)
  if (!ctx) {
    throw new Error(
      `${componentName} must be used within <NestedCheckbox>\u2063`,
    )
  }
  return ctx
}

// Provider component
interface NestedCheckboxProps {
  /**
   * Uncontrolled initial default value of selected option identifiers.
   */
  defaultValue?: string[]
  /**
   * Controlled value of selected option identifiers.
   */
  value?: string[]
  /**
   * Callback fired whenever the selected identifiers change.
   */
  onValueChange?: (value: string[]) => void
  /**
   * Propagate disabled state to every checkbox inside.
   */
  disabled?: boolean

  children: React.ReactNode

  className?: string
}

export const NestedCheckbox = ({
  value,
  defaultValue = [],
  onValueChange,
  disabled,
  children,
  className,
}: NestedCheckboxProps) => {
  const isControlled = value !== undefined
  const [internalSelected, setInternalSelected] = React.useState<Set<string>>(
    () => new Set(defaultValue),
  )
  const [optionIds, setOptionIds] = React.useState<Set<string>>(new Set())

  const selectedSet = React.useMemo<Set<string>>(
    () => (isControlled ? new Set(value) : internalSelected),
    [isControlled, value, internalSelected],
  )

  const updateSelected = React.useCallback(
    (next: Set<string>) => {
      if (!isControlled) {
        setInternalSelected(new Set(next))
      }
      onValueChange?.(Array.from(next))
    },
    [isControlled, onValueChange],
  )

  const registerOption = React.useCallback((id: string) => {
    setOptionIds((prev) => {
      if (prev.has(id)) return prev
      const next = new Set(prev)
      next.add(id)
      return next
    })
  }, [])

  // Keep a ref of the latest selected set so stable callbacks can access up-to-date value.
  const selectedSetRef = React.useRef<Set<string>>(selectedSet)
  React.useEffect(() => {
    selectedSetRef.current = selectedSet
  }, [selectedSet])

  const unregisterOption = React.useCallback(
    (id: string) => {
      // Remove option id from the list of registered options
      setOptionIds((prev) => {
        if (!prev.has(id)) return prev
        const next = new Set(prev)
        next.delete(id)
        return next
      })

      // Remove option id from selected set using functional update to avoid stale closures
      updateSelected(
        new Set(
          Array.from(selectedSet).filter((selectedId) => selectedId !== id),
        ),
      )
    },
    [updateSelected],
  )

  const toggleOption = React.useCallback(
    (id: string, checked: boolean) => {
      const next = new Set(selectedSet)
      if (checked) {
        next.add(id)
      } else {
        next.delete(id)
      }
      updateSelected(next)
    },
    [selectedSet, updateSelected],
  )

  const selectAll = React.useCallback(() => {
    updateSelected(new Set(optionIds))
  }, [optionIds, updateSelected])

  const deselectAll = React.useCallback(() => {
    updateSelected(new Set())
  }, [updateSelected])

  const contextValue = React.useMemo<NestedCheckboxContextValue>(
    () => ({
      selectedIds: Array.from(selectedSet),
      optionIds: Array.from(optionIds),
      registerOption,
      unregisterOption,
      toggleOption,
      selectAll,
      deselectAll,
      disabled,
    }),
    [
      selectedSet,
      optionIds,
      registerOption,
      unregisterOption,
      toggleOption,
      selectAll,
      deselectAll,
      disabled,
    ],
  )

  return (
    <NestedCheckboxContext.Provider value={contextValue}>
      <div className={cn('flex flex-col gap-2', className)}>{children}</div>
    </NestedCheckboxContext.Provider>
  )
}

NestedCheckbox.displayName = 'NestedCheckbox'

// Option component
interface NestedCheckboxOptionProps
  extends Omit<
    React.ComponentPropsWithoutRef<typeof LabeledCheckbox>,
    'checked' | 'onCheckedChange'
  > {
  /** A unique identifier for the option */
  id: string
}

export const NestedCheckboxOption = ({
  id,
  children,
  className,
  ...props
}: NestedCheckboxOptionProps) => {
  const {
    selectedIds,
    toggleOption,
    registerOption,
    unregisterOption,
    disabled,
  } = useNestedCheckboxContext('NestedCheckboxOption')

  const checked = selectedIds.includes(id)

  React.useEffect(() => {
    registerOption(id)
    return () => unregisterOption(id)
  }, [id, registerOption, unregisterOption])

  return (
    <LabeledCheckbox
      {...props}
      id={id}
      className={cn('ml-10', className)}
      disabled={disabled || props.disabled}
      checked={checked}
      onCheckedChange={(state) => {
        const isChecked = state === true
        toggleOption(id, isChecked)
      }}
    >
      {children}
    </LabeledCheckbox>
  )
}

NestedCheckboxOption.displayName = 'NestedCheckboxOption'

// Parent component
type NestedCheckboxToggleProps = Omit<
  React.ComponentPropsWithoutRef<typeof LabeledCheckbox>,
  'checked' | 'onCheckedChange'
>

export const NestedCheckboxToggle = (props: NestedCheckboxToggleProps) => {
  const { selectedIds, optionIds, selectAll, deselectAll, disabled } =
    useNestedCheckboxContext('NestedCheckboxToggle')

  const totalCount = optionIds.length
  const selectedCount = selectedIds.length

  let checked: boolean | 'indeterminate' = false
  if (selectedCount === 0) {
    checked = false
  } else if (selectedCount === totalCount) {
    checked = true
  } else {
    checked = 'indeterminate'
  }

  return (
    <LabeledCheckbox
      {...props}
      disabled={disabled || props.disabled}
      checked={checked}
      onCheckedChange={(state) => {
        if (state === true) {
          selectAll()
        } else {
          deselectAll()
        }
      }}
    >
      {props.children}
    </LabeledCheckbox>
  )
}

NestedCheckboxToggle.displayName = 'NestedCheckboxToggle'

// Exported types
export type { NestedCheckboxProps }