完善选择地区的树组件

src/components/ui/combobox.tsx

@@ -1,7 +1,7 @@
 'use client'
 
 import * as React from 'react'
-import {CheckIcon, ChevronsUpDown} from 'lucide-react'
+import {CheckIcon, ChevronDown, ChevronRight, ChevronsUpDown} from 'lucide-react'
 
 import {merge} from '@/lib/utils'
 import {Button} from '@/components/ui/button'
@@ -10,7 +10,7 @@ import {
   PopoverContent,
   PopoverTrigger,
 } from '@/components/ui/popover'
-import {ReactNode, useRef, useState} from 'react'
+import {ReactNode, useState} from 'react'
 import {Input} from '@/components/ui/input'
 
 type ComboboxItem = {
@@ -28,6 +28,7 @@ type ComboboxProps = {
   value?: string[]
   onChange?: (value: string[]) => void
   children?: React.ReactNode
+  searchPlaceholder?: string
 }
 
 export function Combobox(props: ComboboxProps) {
@@ -56,22 +57,51 @@ export function Combobox(props: ComboboxProps) {
   const [wait, setWait] = useState(false)
   const [filter, setFilter] = useState<string>('')
   const [filtered, setFiltered] = useState<ComboboxItem[]>([])
+  const [expandedKeys, setExpandedKeys] = useState<Set<string>>(new Set())
 
-  const onFilter = async () => {
+  const toggleExpand = (key: string) => {
+    setExpandedKeys((prev) => {
+      const next = new Set(prev)
+      if (next.has(key)) {
+        next.delete(key)
+      }
+      else {
+        next.add(key)
+      }
+      return next
+    })
+  }
+
+  const onSearch = async () => {
     if (wait) return
     const cond = filter.trim()
-    console.log('onFilter', cond)
     setWait(true)
     if (cond.length > 0) {
-      setFiltered(mapFilter(JSON.parse(JSON.stringify(props.options)), cond))
+      const result = mapFilter(JSON.parse(JSON.stringify(props.options)), cond)
+      setFiltered(result)
+      setExpandedKeys(collectAllKeys(result))
     }
     else {
       setFiltered(props.options)
+      setExpandedKeys(new Set())
     }
-    console.log('onFilter end')
     setWait(false)
   }
 
+  const collectAllKeys = (items: ComboboxItem[]): Set<string> => {
+    const keys = new Set<string>()
+    const walk = (list: ComboboxItem[]) => {
+      list.forEach((item) => {
+        if (item.children?.length) {
+          keys.add(item.value)
+          walk(item.children)
+        }
+      })
+    }
+    walk(items)
+    return keys
+  }
+
   const mapFilter = (items: ComboboxItem[], cond: string): ComboboxItem[] => {
     const nItems: ComboboxItem[] = []
     items.forEach((item) => {
@@ -97,6 +127,7 @@ export function Combobox(props: ComboboxProps) {
         if (status) {
           setFiltered(props.options)
           setFilter('')
+          setExpandedKeys(new Set())
         }
       }}>
       <PopoverTrigger asChild>
@@ -117,18 +148,18 @@ export function Combobox(props: ComboboxProps) {
         </Button>
       </PopoverTrigger>
       <PopoverContent
-        className="p-0 rounded-lg w-[var(--radix-popover-trigger-width)] h-[var(--radix-popover-content-available-height)] flex flex-col overflow-hidden"
+        className="p-0 rounded-lg w-(--radix-popover-trigger-width) h-(--radix-popover-content-available-height) flex flex-col overflow-hidden"
         align="start"
         collisionPadding={6}
       >
         <div className="p-2 flex gap-2 flex-none">
           <Input
             className="h-9 placeholder:text-weak placeholder:text-sm"
-            placeholder="搜索地区"
+            placeholder={props.searchPlaceholder || '搜索地区'}
             value={filter}
             onChange={event => setFilter(event.target.value)}
           />
-          <Button className="h-9" onClick={onFilter} disabled={wait}>
+          <Button className="h-9" onClick={onSearch} disabled={wait}>
             搜索
           </Button>
         </div>
@@ -136,11 +167,13 @@ export function Combobox(props: ComboboxProps) {
           <OptionList
             options={filtered}
             value={values}
-            onChange={(value) => {
-              console.log(value.map(item => item.value))
-              props.onChange?.(value.map(item => item.value))
+            expandedKeys={expandedKeys}
+            onToggleExpand={toggleExpand}
+            onChange={(pathValue) => {
+              props.onChange?.(pathValue)
               setOpen(false)
-            }}/>
+            }}
+          />
         </div>
       </PopoverContent>
     </Popover>
@@ -150,51 +183,76 @@ export function Combobox(props: ComboboxProps) {
 function OptionList(props: {
   options: ComboboxItem[]
   value: string[]
-  onChange: (value: ComboboxItem[]) => void
-  path?: ComboboxItem[]
+  expandedKeys: Set<string>
+  onToggleExpand: (key: string) => void
+  onChange: (value: string[]) => void
+  path?: string[]
   depth?: number
 }) {
   const depth = props.depth || 0
-  const parent = props.path || []
-  const indent = depth * 16
+  const parentPath = props.path || []
 
   return (
-    <ul style={{
-      marginLeft: `${indent}px`,
-    }}>
+    <ul>
       {props.options.map((item, i) => {
-        const path = [...parent, item]
-        const pathValue = path.map(item => item.value)
-        const equal = pathValue.join(`.`) === props.value?.join('.')
+        const path = [...parentPath, item.value]
+        const hasChildren = !!item.children?.length
+        const expanded = props.expandedKeys.has(item.value)
+        const isCurrentLevel = path.length === props.value.filter(Boolean).length
+        const isActivePath = path.every((v, idx) => v === props.value[idx])
+
         return (
           <li key={i}>
-            <OptionItem key={`${i}`} item={item} active={equal} onChange={() => props.onChange?.(path)}/>
-            {item.children?.length
-              && <OptionList depth={depth + 1} options={item.children} value={props.value} path={path} onChange={props.onChange}/>
+            <div
+              className={merge(
+                `transition-colors duration-100 ease-in-out`,
+                `pr-4 py-2 rounded-md`,
+                `flex items-center gap-1`,
+                `hover:bg-muted hover:text-foreground cursor-pointer`,
+                isActivePath ? 'text-foreground' : 'text-muted-foreground',
+              )}
+              style={{paddingLeft: `${depth * 20 + 16}px`}}
+            >
+              {hasChildren ? (
+                <span
+                  className="flex-none cursor-pointer p-0.5 hover:bg-muted rounded shrink-0"
+                  onClick={(e) => {
+                    e.stopPropagation()
+                    props.onToggleExpand(item.value)
+                  }}
+                >
+                  {expanded
+                    ? <ChevronDown className="h-4 w-4"/>
+                    : <ChevronRight className="h-4 w-4"/>
                   }
+                </span>
+              ) : (
+                <span className="w-5 shrink-0"/>
+              )}
+              <span
+                className="flex-1 min-w-0 truncate"
+                onClick={() => props.onChange(path)}
+              >
+                {item.label}
+              </span>
+              {isCurrentLevel && isActivePath && (
+                <CheckIcon className="h-4 w-4 shrink-0"/>
+              )}
+            </div>
+            {hasChildren && expanded && (
+              <OptionList
+                depth={depth + 1}
+                options={item.children!}
+                value={props.value}
+                expandedKeys={props.expandedKeys}
+                onToggleExpand={props.onToggleExpand}
+                onChange={props.onChange}
+                path={path}
+              />
+            )}
           </li>
         )
       })}
     </ul>
   )
 }
-
-function OptionItem(props: {
-  key: string
-  item: ComboboxItem
-  active: boolean
-  onChange?: () => void
-}) {
-  return (
-    <div
-      className={merge(
-        `transition-colors, duration-100 ease-in-out`,
-        `px-4 py-2 text-muted-foreground rounded-md`,
-        `flex justify-between items-center`,
-        `hover:bg-muted hover:text-foreground`,
-      )}
-      onClick={props.onChange}>
-      {props.item.label}
-    </div>
-  )
-}