import type { ReactNode } from 'react'
import { useMemo, useState, useCallback, useReducer } from 'react'

import type { Dict } from '../utils'
import { createContext } from '../utils'

export interface SidemenuContextValue extends Dict {
  // Existing properties
  variant?: string
  size?: string
  isMenuExpanded?: boolean

  // Submenu state management
  submenuStates?: Record<string, boolean>
  setSubmenuState?: (itemId: string, isOpen: boolean) => void
  toggleSubmenu?: (itemId: string) => void

  // Expansion state management
  setIsExpanded?: (isExpanded: boolean) => void
  onExpandedChange?: (isExpanded: boolean) => void

  // Memoized reset functions
  resetAllStates?: () => void
  resetExpansionOnly?: () => void
  resetSubmenuOnly?: () => void
}

// Submenu state reducer for more stable state management
type SubmenuAction =
  | { type: 'SET_SUBMENU'; itemId: string; isOpen: boolean }
  | { type: 'TOGGLE_SUBMENU'; itemId: string }
  | { type: 'RESET_ALL' }
  | { type: 'RESTORE_STATE'; state: Record<string, boolean> }

const submenuReducer = (state: Record<string, boolean>, action: SubmenuAction): Record<string, boolean> => {
  let newState: Record<string, boolean>

  switch (action.type) {
    case 'SET_SUBMENU': {
      if (state[action.itemId] === action.isOpen) {
        return state // Return same reference if no change
      }
      newState = { ...state, [action.itemId]: action.isOpen }
      break
    }
    case 'TOGGLE_SUBMENU': {
      const currentValue = state[action.itemId] || false
      newState = { ...state, [action.itemId]: !currentValue }
      break
    }
    case 'RESTORE_STATE': {
      newState = action.state
      break
    }
    case 'RESET_ALL': {
      newState = {}
      break
    }
    default:
      return state
  }

  // Submenu states are now persisted in context only
  return newState
}

interface SidemenuStateProviderProps {
  children: ReactNode
  isExpandedInitial?: boolean
  onExpandedChange?: (isExpanded: boolean) => void
  variant?: string
  size?: string
  isDark?: boolean
}

const [SidemenuProvider, useSidemenuContext] = createContext<SidemenuContextValue>({ isOptional: true })

// Enhanced provider that manages state
export const SidemenuStateProvider = ({
  children,
  isExpandedInitial = true,
  onExpandedChange,
  variant,
  size,
  isDark,
}: SidemenuStateProviderProps) => {
  // Pure uncontrolled state without persistence - context only
  const [isExpanded, setIsExpanded] = useState(isExpandedInitial)

  // Initialize submenu states as empty - persisted in context only
  const [submenuStates, dispatchSubmenu] = useReducer(submenuReducer, {})

  // Submenu state management functions using reducer dispatch - memoized
  const setSubmenuState = useCallback((itemId: string, isOpen: boolean) => {
    dispatchSubmenu({ type: 'SET_SUBMENU', itemId, isOpen })
  }, [])

  const toggleSubmenu = useCallback((itemId: string) => {
    dispatchSubmenu({ type: 'TOGGLE_SUBMENU', itemId })
  }, [])

  // Enhanced setIsExpanded that handles callbacks - memoized
  const handleSetIsExpanded = useCallback(
    (newIsExpanded: boolean) => {
      setIsExpanded(newIsExpanded)
      onExpandedChange?.(newIsExpanded)
    },
    [onExpandedChange],
  )

  // Memoized reset functions for external use
  const resetAllStates = useCallback(() => {
    dispatchSubmenu({ type: 'RESET_ALL' })
    const defaultExpanded = isExpandedInitial
    setIsExpanded(defaultExpanded)
    onExpandedChange?.(defaultExpanded)
  }, [isExpandedInitial, onExpandedChange])

  const resetExpansionOnly = useCallback(() => {
    const defaultExpanded = isExpandedInitial
    setIsExpanded(defaultExpanded)
    onExpandedChange?.(defaultExpanded)
  }, [isExpandedInitial, onExpandedChange])

  const resetSubmenuOnly = useCallback(() => {
    dispatchSubmenu({ type: 'RESET_ALL' })
  }, [])

  // Memoized context value with all state management functions
  const contextValue = useMemo(
    () => ({
      size,
      variant,
      isDark,
      isMenuExpanded: isExpanded,
      submenuStates,
      setSubmenuState,
      toggleSubmenu,
      setIsExpanded: handleSetIsExpanded,
      onExpandedChange,
      resetAllStates,
      resetExpansionOnly,
      resetSubmenuOnly,
    }),
    [
      size,
      variant,
      isDark,
      isExpanded,
      submenuStates,
      setSubmenuState,
      toggleSubmenu,
      handleSetIsExpanded,
      onExpandedChange,
      resetAllStates,
      resetExpansionOnly,
      resetSubmenuOnly,
    ],
  )

  return <SidemenuProvider value={contextValue}>{children}</SidemenuProvider>
}

export { SidemenuProvider, useSidemenuContext }
