import {
  forwardRef,
  useCallback,
  useEffect,
  useImperativeHandle,
  useMemo,
  useRef,
  useState,
} from 'react'

import {Stack} from '../../primitives'
import {_findNextItemElement, _findPrevItemElement, _focusItemElement} from './helpers'
import {TreeContext} from './treeContext'
import {TreeContextValue, TreeState} from './types'

/**
 * This API might change. DO NOT USE IN PRODUCTION.
 * @beta
 */
export interface TreeProps {
  gap?: number | number[]
  /**
   * @deprecated Use `gap` instead. `space` will be removed in v4.
   */
  space?: number | number[]
}

/**
 * This API might change. DO NOT USE IN PRODUCTION.
 * @beta
 */
export const Tree = forwardRef(function Tree(
  props: TreeProps &
    Omit<React.HTMLProps<HTMLDivElement>, 'align' | 'as' | 'height' | 'ref' | 'role' | 'wrap'>,
  forwardedRef: React.ForwardedRef<HTMLDivElement>,
): React.JSX.Element {
  const {children, gap, space: deprecated_space = 1, onFocus, ...restProps} = props
  const spacing = gap === undefined ? deprecated_space : gap
  const ref = useRef<HTMLDivElement | null>(null)
  const [focusedElement, setFocusedElement] = useState<HTMLElement | null>(null)
  const focusedElementRef = useRef(focusedElement)
  const path: string[] = useMemo(() => [], [])
  const [itemElements, setItemElements] = useState<HTMLElement[]>([])
  const [state, setState] = useState<TreeState>({})
  const stateRef = useRef(state)

  useImperativeHandle<HTMLDivElement | null, HTMLDivElement | null>(forwardedRef, () => ref.current)

  useEffect(() => {
    focusedElementRef.current = focusedElement
  }, [focusedElement])

  useEffect(() => {
    stateRef.current = state
  }, [state])

  const registerItem = useCallback(
    (element: HTMLElement, path: string, expanded: boolean, selected: boolean) => {
      setState((s) => ({...s, [path]: {element, expanded}}))

      if (selected) {
        setFocusedElement(element)
      }

      return () => {
        setState((s) => {
          const newState = {...s}

          delete newState[path]

          return newState
        })
      }
    },
    [],
  )

  const setExpanded = useCallback((path: string, expanded: boolean) => {
    setState((s) => {
      const itemState = s[path]

      if (!itemState) return s

      return {...s, [path]: {...itemState, expanded}}
    })
  }, [])

  const contextValue: TreeContextValue = useMemo(
    () => ({
      version: 0.0,
      focusedElement: focusedElement || itemElements[0] || null,
      level: 0,
      path,
      registerItem,
      setExpanded,
      setFocusedElement,
      gap: spacing,
      space: spacing,
      state,
    }),
    [focusedElement, itemElements, path, registerItem, setExpanded, spacing, state],
  )

  const handleKeyDown = useCallback(
    (event: React.KeyboardEvent<HTMLDivElement>) => {
      if (!focusedElementRef.current) return

      if (event.key === 'ArrowDown') {
        event.preventDefault()

        const nextEl = _findNextItemElement(
          stateRef.current,
          itemElements,
          focusedElementRef.current,
        )

        if (nextEl) {
          _focusItemElement(nextEl)
          setFocusedElement(nextEl)
        }

        return
      }

      if (event.key === 'ArrowUp') {
        event.preventDefault()

        const prevEl = _findPrevItemElement(
          stateRef.current,
          itemElements,
          focusedElementRef.current,
        )

        if (prevEl) {
          _focusItemElement(prevEl)
          setFocusedElement(prevEl)
        }

        return
      }

      if (event.key === 'ArrowLeft') {
        event.preventDefault()

        const itemKey = focusedElementRef.current.getAttribute('data-tree-key')

        if (!itemKey) return

        const itemState = stateRef.current[itemKey]

        if (!itemState) return

        if (itemState.expanded) {
          setState((s) => {
            const itemState = s[itemKey]

            if (!itemState) return s

            return {...s, [itemKey]: {...itemState, expanded: false}}
          })
        } else {
          const itemPath = itemKey.split('/')

          itemPath.pop()

          const parentKey = itemPath.join('/')
          const parentState = parentKey && stateRef.current[parentKey]

          if (parentState) {
            parentState.element.focus()
            setFocusedElement(parentState.element)
          }
        }

        return
      }

      if (event.key === 'ArrowRight') {
        event.preventDefault()

        const focusedKey = focusedElementRef.current.getAttribute('data-tree-key')

        if (!focusedKey) return

        if (!stateRef.current[focusedKey]?.expanded) {
          setState((s) => {
            const itemState = s[focusedKey]

            if (!itemState) return s

            return {...s, [focusedKey]: {...itemState, expanded: true}}
          })
        }

        return
      }
    },
    [itemElements],
  )

  const handleFocus = useCallback(
    (event: React.FocusEvent<HTMLDivElement>) => {
      setFocusedElement(event.target)

      // Call the element's `focus` handler
      onFocus?.(event)
    },
    [onFocus],
  )

  useEffect(() => {
    if (!ref.current) return
    const _itemElements = Array.from(
      ref.current.querySelectorAll('[data-ui="TreeItem"]'),
    ) as HTMLElement[]

    setItemElements(_itemElements)
  }, [children])

  return (
    <TreeContext.Provider value={contextValue}>
      <Stack
        as="ul"
        data-ui="Tree"
        {...restProps}
        onFocus={handleFocus}
        onKeyDown={handleKeyDown}
        ref={ref}
        role="tree"
        gap={spacing}
      >
        {children}
      </Stack>
    </TreeContext.Provider>
  )
})
Tree.displayName = 'ForwardRef(Tree)'