ci4sManagement-cloud/react-ui/src/components/ResourceSelectorModal/index.tsx

/*
 * @Author: 赵伟
 * @Date: 2024-04-11 16:31:18
 * @Description: 选择数据集、模型、镜像
 */

import KFIcon from '@/components/KFIcon';
import KFModal from '@/components/KFModal';
import { CommonTabKeys } from '@/enums';
import { ResourceFileData } from '@/pages/Dataset/config';
import { to } from '@/utils/promise';
import type { GetRef, ModalProps, TreeDataNode, TreeProps } from 'antd';
import { Input, Tabs, Tree } from 'antd';
import React, { useEffect, useMemo, useRef, useState } from 'react';
import { ResourceSelectorType, selectorTypeConfig } from './config';
import styles from './index.less';
export { ResourceSelectorType, selectorTypeConfig };

// 选择数据集、模型、镜像的返回类型
export type ResourceSelectorResponse = {
  activeTab: CommonTabKeys; // 是我的还是公开的
  id: string; // 数据集\模型\镜像 id
  name: string; // 数据集\模型\镜像 name
  version: string; // 数据集\模型\镜像版本
  path: string; // 数据集\模型\镜像版本路径
  identifier: string; // 数据集\模型 identifier，镜像这个字段为空
  owner: string; // 数据集\模型 owner，镜像这个字段为空
};

export interface ResourceSelectorModalProps extends Omit<ModalProps, 'onOk'> {
  /** 类型，数据集、模型、镜像 */
  type: ResourceSelectorType;
  /** 默认展开的节点 */
  defaultExpandedKeys?: React.Key[];
  /** 默认展开的节点 */
  defaultCheckedKeys?: React.Key[];
  /** 默认激活的 Tab */
  defaultActiveTab?: CommonTabKeys;
  /**
   * 确认回调
   * @param params 选择的数据
   */
  onOk?: (params: ResourceSelectorResponse | undefined) => void;
}

type TreeRef = GetRef<typeof Tree<TreeDataNode>>;

// 更新树形结构的 children
const updateChildren = (parentId: string, children: TreeDataNode[]) => {
  return (node: TreeDataNode) => {
    if (node.key === parentId) {
      return {
        ...node,
        children,
      };
    }
    return node;
  };
};

// 得到数据集\模型\镜像 id 和下属版本号
const getIdAndVersion = (versionKey: string) => {
  const index = versionKey.indexOf('-');
  const id = versionKey.slice(0, index);
  const version = versionKey.slice(index + 1);
  return {
    id,
    version,
  };
};

/** 选择数据集、模型、镜像的弹框，推荐使用函数的方式打开 */
function ResourceSelectorModal({
  type,
  defaultExpandedKeys = [],
  defaultCheckedKeys = [],
  defaultActiveTab = CommonTabKeys.Private,
  onOk,
  ...rest
}: ResourceSelectorModalProps) {
  const [activeTab, setActiveTab] = useState<CommonTabKeys>(defaultActiveTab);
  const [expandedKeys, setExpandedKeys] = useState<React.Key[]>([]);
  const [checkedKeys, setCheckedKeys] = useState<React.Key[]>([]);
  const [loadedKeys, setLoadedKeys] = useState<React.Key[]>([]);
  const [originTreeData, setOriginTreeData] = useState<TreeDataNode[]>([]);
  const [files, setFiles] = useState<ResourceFileData[]>([]);
  const [versionPath, setVersionPath] = useState('');
  const [searchText, setSearchText] = useState('');
  const [firstLoadList, setFirstLoadList] = useState(false);
  const [firstLoadVersions, setFirstLoadVersions] = useState(false);
  const treeRef = useRef<TreeRef>(null);
  const config = selectorTypeConfig[type];

  useEffect(() => {
    setExpandedKeys([]);
    setCheckedKeys([]);
    setLoadedKeys([]);
    setFiles([]);
    setVersionPath('');
    setSearchText('');
    getTreeData();
  }, [activeTab, type]);

  const treeData = useMemo(
    () =>
      originTreeData.filter((v) =>
        (v.title as string).toLowerCase()?.includes(searchText.toLowerCase()),
      ),
    [originTreeData, searchText],
  );

  // 获取数据集\模型\镜像列表
  const getTreeData = async () => {
    const isPublic = activeTab === CommonTabKeys.Private ? false : true;
    const [res] = await to(config.getList(isPublic));
    if (res) {
      setOriginTreeData(res);

      // 恢复上一次的 Expand 操作
      restoreLastExpand();
    } else {
      setOriginTreeData([]);
    }
  };

  // 获取数据集\模型\镜像版本列表
  const getVersions = async (parentId: string, parentNode: any) => {
    const [res, error] = await to(config.getVersions(parentId, parentNode));
    if (res) {
      // 更新 treeData children
      setOriginTreeData((prev) => prev.map(updateChildren(parentId, res)));

      // 缓存 loadedKeys
      const index = loadedKeys.find((v) => v === parentId);
      if (!index) {
        setLoadedKeys((prev) => prev.concat(parentId));
      }

      // 恢复上一次的 Check 操作
      setTimeout(() => {
        restoreLastCheck(parentId, res);
      }, 300);
    } else {
      setExpandedKeys([]);
      return Promise.reject(error);
    }
  };

  // 获取版本下的文件
  const getFiles = async (parentId: string, parentNode: any) => {
    const [res] = await to(config.getFiles(parentId, parentNode));
    if (res) {
      setVersionPath(res.path);
      setFiles(res.content);
    } else {
      setVersionPath('');
      setFiles([]);
    }
  };

  // 动态加载 tree children
  const onLoadData = ({ key, children, ...rest }: TreeDataNode) => {
    if (children) {
      return Promise.resolve();
    } else {
      return getVersions(key as string, rest);
    }
  };

  // 扩展
  const onExpand: TreeProps['onExpand'] = (expandedKeysValue) => {
    const lastKeys = expandedKeysValue.slice(-1);
    setExpandedKeys(lastKeys);
  };

  // 选中
  const onCheck: TreeProps['onCheck'] = (checkedKeysValue, { checkedNodes }) => {
    const lastKeys = (checkedKeysValue as React.Key[]).slice(-1);
    setCheckedKeys(lastKeys);
    if (lastKeys.length && checkedNodes.length) {
      const last = lastKeys[0] as string;
      const lastNode = checkedNodes[checkedNodes.length - 1];
      getFiles(last, lastNode);
    } else {
      setVersionPath('');
      setFiles([]);
    }
  };

  // 恢复上一次的 Expand 操作
  // 判断是否有 defaultExpandedKeys，如果有，设置 expandedKeys
  // fisrtLoadList 标志位
  const restoreLastExpand = () => {
    if (!firstLoadList && defaultExpandedKeys.length > 0) {
      setTimeout(() => {
        setExpandedKeys(defaultExpandedKeys);
        setFirstLoadList(true);
        setTimeout(() => {
          treeRef.current?.scrollTo({ key: defaultExpandedKeys[0], align: 'bottom' });
        }, 100);
      }, 0);
    }
  };

  // 恢复上一次的 Check 操作
  // 判断是否有 defaultCheckedKeys，如果有，设置 checkedKeys，并且调用获取文件列表接口
  // fisrtLoadVersions 标志位
  const restoreLastCheck = (parentId: string, versions: TreeDataNode[]) => {
    if (!firstLoadVersions && defaultCheckedKeys.length > 0) {
      const last = defaultCheckedKeys[0] as string;
      const { id } = getIdAndVersion(last);
      // 判断正在打开的 id 和 defaultCheckedKeys 的 id 是否一致
      if (id === parentId) {
        setTimeout(() => {
          setCheckedKeys(defaultCheckedKeys);
          const parentNode = versions.find((v) => v.key === last);
          getFiles(last, parentNode);
          setFirstLoadVersions(true);
          setTimeout(() => {
            treeRef?.current?.scrollTo({
              key: defaultCheckedKeys[0],
              align: 'bottom',
            });
          }, 100);
        }, 0);
      }
    }
  };

  // 提交
  const handleOk = () => {
    if (checkedKeys.length > 0) {
      const last = checkedKeys[0] as string;
      const { id, version } = getIdAndVersion(last);
      const treeNode = treeData.find((v) => v.key === id) as any;
      const name = (treeNode?.title ?? '') as string;
      const identifier = (treeNode?.identifier ?? '') as string;
      const owner = (treeNode?.owner ?? '') as string;
      const res = {
        id,
        name,
        path: versionPath,
        version,
        identifier,
        owner,
        activeTab: activeTab,
      };
      onOk?.(res);
    } else {
      onOk?.(undefined);
    }
  };

  const title = `选择${config.name}`;
  const palceholder = `请输入${config.name}名称`;
  const fileTitle =
    type === ResourceSelectorType.Mirror ? '已选镜像' : `已选${config.name}文件（${files.length}）`;
  const tabItems = config.tabItems;
  const titleImg = config.modalIcon;

  return (
    <KFModal {...rest} title={title} image={titleImg} onOk={handleOk} width={920} destroyOnClose>
      <div>
        <Tabs
          activeKey={activeTab}
          items={tabItems}
          onChange={(e) => setActiveTab(e as CommonTabKeys)}
          className={styles['model-tabs']}
        />
        <div className={styles['model-selector']}>
          <div className={styles['model-selector__left']}>
            <Input
              className={styles['model-selector__left__search']}
              placeholder={palceholder}
              allowClear
              variant="borderless"
              value={searchText}
              onChange={(e) => setSearchText(e.target.value)}
              prefix={
                <KFIcon
                  type="icon-sousuo"
                  color="rgba(22,100,255,0.4)"
                  style={{ height: '15px' }}
                />
              }
              // prefix={<Icon icon="local:magnifying-glass" style={{ height: '15px' }} />}
            />
            <Tree
              ref={treeRef}
              rootStyle={{ backgroundColor: 'transparent' }}
              loadData={onLoadData}
              treeData={treeData}
              onCheck={onCheck}
              checkedKeys={checkedKeys}
              multiple={false}
              selectable={false}
              height={324}
              loadedKeys={loadedKeys}
              expandedKeys={expandedKeys}
              onExpand={onExpand}
              checkable
              titleRender={(nodeData) => {
                return (
                  <span
                    className={styles['model-selector__left__tree-title']}
                    style={{ width: nodeData.isLeaf ? '370px' : '420px' }}
                  >
                    {nodeData.title as string}
                  </span>
                );
              }}
            />
          </div>
          <div className={styles['model-selector__right']}>
            <div className={styles['model-selector__right__title']}>{fileTitle}</div>
            <div className={styles['model-selector__right__files']}>
              {files.map((v) => (
                <div key={v.url} className={styles['model-selector__right__files__file']}>
                  {v.file_name}
                </div>
              ))}
            </div>
          </div>
        </div>
      </div>
    </KFModal>
  );
}

export default ResourceSelectorModal;