labelmain-demo/components/tree-select/tree.tsx

import React, { useEffect, useMemo, useState } from "react"
import {
  Popover,
  Text,
  UnstyledButton,
  Input,
  Checkbox,
  useTree,
  getTreeExpandedState,
  RenderTreeNodePayload,
  Flex,
  TreeNodeData,
  Combobox,
  CloseButton,
  Stack,
  TextInput,
  InputProps,
} from "@mantine/core"
import { useDisclosure } from "@mantine/hooks"
import { IconChevronUp, IconSearch } from "@tabler/icons-react"
import { VirtualizedTree } from "./VirtualizedTree"
import {
  filterTreeData_dfs,
  findLabel,
  getStateFromDataByName,
} from "./libs/util"

interface VirtualTreeSelectProps extends InputProps {
  data: TreeNodeData[]
  value: string
  onChange: (value: string) => void
  placeholder?: string
  setCurrCityCode: (code: string) => void
}

export function VirtualTreeSelect({
  data,
  value,
  onChange,
  setCurrCityCode,
  ...props
}: VirtualTreeSelectProps) {
  const [opened, { toggle }] = useDisclosure(false)
  const [searchName, setSearchName] = useState<string>("")

  const filteredTreeData = useMemo(() => {
    return filterTreeData_dfs(data, searchName)
  }, [searchName, data])

  const tree = useTree({
    initialExpandedState: getTreeExpandedState(filteredTreeData, []),
    initialCheckedState: [],
  })

  useEffect(() => {
    if (filteredTreeData?.length && searchName) {
      let state = getStateFromDataByName(filteredTreeData)
      tree.setExpandedState(state)
    }
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [filteredTreeData, searchName])

  useEffect(() => {
    setCurrCityCode(tree.selectedState[0])
    let label = findLabel(data, tree.selectedState[0])
    onChange(label.path || "")
  }, [data, onChange, setCurrCityCode, tree.selectedState])

  const ChevronIcon = ({
    hasChildren,
    expanded,
  }: {
    hasChildren: boolean
    expanded: boolean
  }) => {
    if (hasChildren) {
      return (
        <IconChevronUp
          size={16}
          stroke={2.5}
          style={{
            transform: expanded ? "rotate(180deg)" : "rotate(90deg)",
          }}
        />
      )
    }

    return null
  }

  const Leaf = React.memo(
    ({
      node,
      expanded,
      elementProps,
      selected,
      tree,
    }: RenderTreeNodePayload) => {
      return (
        <Flex
          {...elementProps}
          onClick={() => {}}
          gap={12}
          px={16}
          justify={"space-between"}
          flex={1}
          align="stretch"
          bdrs={4}
          pl={(node.nodeProps?.level || 1) * 16}
          style={{
            overflowX: "hidden",
            cursor: "auto",
            height: 36,
          }}>
          <Flex align={"center"} gap={12} flex={1} onClick={() => {}}>
            <Checkbox.Indicator
              checked={selected}
              size="xs"
              radius="sm"
              onClick={() => {
                if (!selected) {
                  tree.select(node.value)
                } else {
                  tree.select("")
                }
              }}
            />
            <Text
              lh={"22px"}
              size="15px"
              title={node.label as string}
              style={{
                whiteSpace: "nowrap",
                overflow: "hidden",
                textOverflow: "ellipsis",
              }}>
              {node.label}
            </Text>
          </Flex>
          <Flex gap={24} align={"center"}>
            <UnstyledButton
              h={16}
              style={{
                display: "flex",
              }}
              onClick={() => tree.toggleExpanded(node.value)}>
              <ChevronIcon
                hasChildren={!!(node.children && node.children.length)}
                expanded={expanded}
              />
            </UnstyledButton>
          </Flex>
        </Flex>
      )
    }
  )

  return (
    <Popover
      opened={opened}
      onChange={toggle}
      width="target"
      position="bottom-start"
      shadow="md"
      keepMounted={false} // 性能优化：关闭时销毁 DOM
    >
      <Popover.Target>
        <Input
          component="button"
          pointer
          onClick={toggle}
          rightSection={<Combobox.Chevron />}
          // style={{ width: 250 }}
          {...props}>
          {value || "请选择..."}
        </Input>
      </Popover.Target>
      <Popover.Dropdown p={0}>
        <TextInput
          m={"xs"}
          radius={"md"}
          placeholder="搜索省、市、区名称"
          value={searchName}
          onChange={(e) => {
            setSearchName(e.currentTarget.value)
          }}
          rightSection={
            searchName ? (
              <CloseButton
                size="1rem"
                aria-label="Clear input"
                onClick={() => {
                  setSearchName("")
                }}
                style={{
                  display: searchName ? undefined : "none",
                  cursor: "pointer",
                }}
              />
            ) : (
              <IconSearch size="1rem" />
            )
          }
        />
        {filteredTreeData.length === 0 ? (
          <Stack flex={1} h={"250px"} justify="center" align="center">
            <Text c="dimmed">暂无数据</Text>
          </Stack>
        ) : (
          <>
            {/* 虚拟列表容器：必须有固定高度 */}
            <VirtualizedTree
              height={"450px"}
              tree={tree}
              data={filteredTreeData}
              levelOffset={30}
              renderNode={(payload) => <Leaf {...payload} />}
            />
          </>
        )}
      </Popover.Dropdown>
    </Popover>
  )
}