metabuilder/fakemui/react/components/lab/TreeView.tsx

'use client'

import React, { useState, createContext, useContext, forwardRef } from 'react'
import { classNames } from '../utils/classNames'

interface TreeViewContextType {
  expanded: string[]
  selected: string[]
  multiSelect: boolean
  onNodeToggle: (event: React.SyntheticEvent, nodeId: string) => void
  onNodeSelect: (event: React.SyntheticEvent, nodeId: string) => void
  defaultCollapseIcon?: React.ReactNode
  defaultEndIcon?: React.ReactNode
  defaultExpandIcon?: React.ReactNode
  defaultParentIcon?: React.ReactNode
}

const TreeViewContext = createContext<TreeViewContextType>({
  expanded: [],
  selected: [],
  multiSelect: false,
  onNodeToggle: () => {},
  onNodeSelect: () => {},
})

export interface TreeViewProps extends React.HTMLAttributes<HTMLUListElement> {
  children?: React.ReactNode
  defaultCollapseIcon?: React.ReactNode
  defaultEndIcon?: React.ReactNode
  defaultExpandIcon?: React.ReactNode
  defaultParentIcon?: React.ReactNode
  expanded?: string[]
  defaultExpanded?: string[]
  selected?: string | string[]
  defaultSelected?: string | string[]
  multiSelect?: boolean
  onNodeToggle?: (event: React.SyntheticEvent, nodeIds: string[]) => void
  onNodeSelect?: (event: React.SyntheticEvent, nodeIds: string | string[]) => void
  disableSelection?: boolean
}

export function TreeView({
  children,
  defaultCollapseIcon,
  defaultEndIcon,
  defaultExpandIcon,
  defaultParentIcon,
  expanded: controlledExpanded,
  defaultExpanded = [],
  selected: controlledSelected,
  defaultSelected = [],
  multiSelect = false,
  onNodeToggle,
  onNodeSelect,
  disableSelection = false,
  className,
  ...props
}: TreeViewProps) {
  const [internalExpanded, setInternalExpanded] = useState<string[]>(defaultExpanded)
  const [internalSelected, setInternalSelected] = useState<string[]>(
    Array.isArray(defaultSelected) ? defaultSelected : [defaultSelected]
  )

  const expanded = controlledExpanded ?? internalExpanded
  const selected = controlledSelected ?? internalSelected

  const handleNodeToggle = (event: React.SyntheticEvent, nodeId: string) => {
    const newExpanded = expanded.includes(nodeId) ? expanded.filter((id) => id !== nodeId) : [...expanded, nodeId]

    if (controlledExpanded === undefined) {
      setInternalExpanded(newExpanded)
    }
    onNodeToggle?.(event, newExpanded)
  }

  const handleNodeSelect = (event: React.SyntheticEvent, nodeId: string) => {
    if (disableSelection) return

    let newSelected: string[]
    if (multiSelect) {
      newSelected = (Array.isArray(selected) ? selected : [selected]).includes(nodeId)
        ? (Array.isArray(selected) ? selected : [selected]).filter((id) => id !== nodeId)
        : [...(Array.isArray(selected) ? selected : [selected]), nodeId]
    } else {
      newSelected = [nodeId]
    }

    if (controlledSelected === undefined) {
      setInternalSelected(newSelected)
    }
    onNodeSelect?.(event, multiSelect ? newSelected : nodeId)
  }

  return (
    <TreeViewContext.Provider
      value={{
        expanded,
        selected: Array.isArray(selected) ? selected : [selected],
        multiSelect,
        onNodeToggle: handleNodeToggle,
        onNodeSelect: handleNodeSelect,
        defaultCollapseIcon,
        defaultEndIcon,
        defaultExpandIcon,
        defaultParentIcon,
      }}
    >
      <ul className={classNames('fakemui-tree-view', className)} role="tree" aria-multiselectable={multiSelect} {...props}>
        {children}
      </ul>
    </TreeViewContext.Provider>
  )
}

export interface TreeItemProps extends React.LiHTMLAttributes<HTMLLIElement> {
  nodeId: string
  label?: React.ReactNode
  children?: React.ReactNode
  icon?: React.ReactNode
  expandIcon?: React.ReactNode
  collapseIcon?: React.ReactNode
  endIcon?: React.ReactNode
  disabled?: boolean
  ContentComponent?: React.ElementType
  ContentProps?: any
}

export const TreeItem = forwardRef<HTMLLIElement, TreeItemProps>(function TreeItem(
  { nodeId, label, children, icon, expandIcon, collapseIcon, endIcon, disabled = false, className, ContentComponent, ContentProps = {}, ...props },
  ref
) {
  const { expanded, selected, onNodeToggle, onNodeSelect, defaultCollapseIcon, defaultEndIcon, defaultExpandIcon } =
    useContext(TreeViewContext)

  const isExpanded = expanded.includes(nodeId)
  const isSelected = selected.includes(nodeId)
  const hasChildren = React.Children.count(children) > 0

  const handleClick = (event: React.MouseEvent) => {
    if (disabled) return
    onNodeSelect(event, nodeId)
  }

  const handleExpandClick = (event: React.MouseEvent) => {
    event.stopPropagation()
    if (disabled) return
    onNodeToggle(event, nodeId)
  }

  const displayIcon = hasChildren
    ? isExpanded
      ? collapseIcon || defaultCollapseIcon || '▼'
      : expandIcon || defaultExpandIcon || '▶'
    : endIcon || defaultEndIcon || null

  return (
    <li
      ref={ref}
      className={classNames('fakemui-tree-item', className, {
        'fakemui-tree-item-expanded': isExpanded,
        'fakemui-tree-item-selected': isSelected,
        'fakemui-tree-item-disabled': disabled,
      })}
      role="treeitem"
      aria-expanded={hasChildren ? isExpanded : undefined}
      aria-selected={isSelected}
      aria-disabled={disabled}
      {...props}
    >
      <div className="fakemui-tree-item-content" onClick={handleClick} {...ContentProps}>
        {hasChildren && (
          <span className="fakemui-tree-item-icon-container" onClick={handleExpandClick}>
            {displayIcon}
          </span>
        )}
        {!hasChildren && displayIcon && <span className="fakemui-tree-item-icon-container">{displayIcon}</span>}
        {icon && <span className="fakemui-tree-item-icon">{icon}</span>}
        <span className="fakemui-tree-item-label">{label}</span>
      </div>
      {hasChildren && isExpanded && (
        <ul className="fakemui-tree-item-group" role="group">
          {children}
        </ul>
      )}
    </li>
  )
})

export default TreeView