TreeSelect解决回显问题

import type { GetProp, TreeSelectProps } from 'antd';
import { TreeSelect } from 'antd';
import { cloneDeep as _cloneDeep, omit as _omit, set as _set } from 'lodash-es';
import React, { useEffect, useState, useRef } from 'react';
import newRequest from '../request';
import type { FieldProps } from '../typing';

type DefaultOptionType = GetProp<TreeSelectProps, 'treeData'>[number];

const normalizeNodes = (list: any[] = []): Omit<DefaultOptionType, 'label'>[] =>
  list?.map((item) => ({
    key: item.value,
    value: item.value,
    title: item.label ?? item.value,
    selectable: true,
    isLeaf: item?.hasSubDept !== 'Y',
    children: [],
    ...item,
  }));

const appendChildren = (
  nodes: Omit<DefaultOptionType, 'label'>[],
  targetKey: React.Key,
  children: Omit<DefaultOptionType, 'label'>[],
): Omit<DefaultOptionType, 'label'>[] =>
  nodes.map((node) => {
    if (node.key === targetKey) {
      return {
        ...node,
        children,
      };
    }
    if (node.children && node.children.length) {
      return {
        ...node,
        children: appendChildren(node.children as any, targetKey, children),
      };
    }
    return node;
  });

/**
 * 树形选择器通用组件
 * @param props
 * @constructor
 */
const ComponentTreeSelect: React.FC = (props: FieldProps) => {
  const [treeData, setTreeData] = useState<Omit<DefaultOptionType, 'label'>[]>(
    [],
  );
  // 用于存储 value 到 label 的映射，解决回显问题
  const [valueMap, setValueMap] = useState<Record<string, string>>({});
  const isInitRef = useRef(false);

  /**
   * 获取显示值的函数
   */
  const fetchDisplayValue = async (value: string) => {
    if (!value || valueMap[value]) return;

    const newSearchParams = _cloneDeep(props?.searchParams ?? {});
    _set(newSearchParams, 'data.value', value);

    try {
      const res: any = await newRequest({ cache: true, ...newSearchParams });
      const nodeData = res?.data?.data?.[0];
      if (nodeData) {
        setValueMap((prev) => ({
          ...prev,
          [value]: nodeData.label || nodeData.title || value,
        }));
      }
    } catch (error) {
      console.warn('Failed to fetch display value:', error);
    }
  };

  /**
   * 初始化加载根节点数据
   */
  useEffect(() => {
    const loadRootData = async () => {
      try {
        const rootRequest = newRequest({
          cache: true,
          ...(props?.searchParams ?? {}),
        });

        const rootRes: any = await rootRequest;
        const finalTreeData = normalizeNodes(rootRes?.data?.data || []);
        setTreeData(finalTreeData);
        isInitRef.current = true;

        // 如果有初始值，获取其显示名称
        if (props.value) {
          fetchDisplayValue(props.value);
        }
      } catch (error) {
        console.warn('Failed to load root data:', error);
        isInitRef.current = true;
      }
    };

    loadRootData();
  }, []);

  /**
   * 监听 value 变化，获取对应的显示名称
   */
  useEffect(() => {
    if (isInitRef.current && props.value) {
      fetchDisplayValue(props.value);
    }
  }, [props.value]);

  const onLoadData: TreeSelectProps['loadData'] = ({ key }) =>
    new Promise((resolve) => {
      if (!key) {
        resolve(undefined);
        return;
      }
      const newSearchParams = _cloneDeep(props?.searchParams ?? {});
      _set(newSearchParams, 'data.code', key);
      newRequest({ cache: true, ...newSearchParams }).then((res: any) => {
        const childrenNodes = normalizeNodes(res?.data?.data || []);
        setTreeData((prev) => appendChildren(prev, key, childrenNodes));
        resolve(undefined);
      });
    });

  // 创建增强的 treeData，确保选中值能正确显示
  const enhancedTreeData = React.useMemo(() => {
    // 更新现有节点的显示名称
    const updateNodeTitles = (
      nodes: Omit<DefaultOptionType, 'label'>[],
    ): Omit<DefaultOptionType, 'label'>[] => {
      return nodes.map((node) => {
        const displayName = valueMap[node.value as string];
        return {
          ...node,
          title: displayName || node.title,
          children: node.children ? updateNodeTitles(node.children as any) : [],
        };
      });
    };

    let result = updateNodeTitles(treeData);

    // 如果当前选中的值不在树中，但有对应的显示名称，临时添加一个节点
    if (props.value && valueMap[props.value]) {
      const existsInTree = (nodes: any[], value: string): boolean => {
        return nodes.some(
          (node) =>
            node.value === value ||
            (node.children && existsInTree(node.children, value)),
        );
      };

      if (!existsInTree(result, props.value)) {
        // 添加临时节点到数据开头，确保 TreeSelect 能找到并显示
        // 但在下拉列表中隐藏这个节点
        result = [
          {
            key: props.value,
            value: props.value,
            title: valueMap[props.value],
            selectable: false, // 不可选择，避免用户在下拉中看到
            disabled: true, // 禁用状态
            isLeaf: true,
            children: [],
            style: { display: 'none' }, // 在下拉列表中隐藏
            className: 'tree-select-hidden-node', // 添加类名便于识别
          },
          ...result,
        ];
      }
    }

    return result;
  }, [treeData, valueMap, props.value]);

  return (
    <TreeSelect
      maxTagCount={20}
      treeData={enhancedTreeData}
      loadData={onLoadData}
      {..._omit(props, ['searchParams'])}
    />
  );
};

export default ComponentTreeSelect;