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 { type MirrorVersionData } from '@/pages/Mirror/Info';
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
  identifier: string; // 数据集\模型 identifier
  owner: string; // 数据集\模型 owner
  version: string; // 数据集\模型 version
  path: string; // 数据集\模型 版本路径
} & MirrorVersionData;

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 [versionDesc, setVersionDesc] = useState<string | undefined>(undefined);
  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];

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

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

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

    setExpandedKeys([]);
    setCheckedKeys([]);
    setLoadedKeys([]);
    setFiles([]);
    setVersionDesc(undefined);
    setVersionPath('');
    setSearchText('');
    getTreeData();
  }, [activeTab, config]);

  useEffect(() => {
    // 恢复上一次的 Expand 操作
    // 判断是否有 defaultExpandedKeys，如果有，设置 expandedKeys
    // fisrtLoadList 标志位
    const restoreLastExpand = () => {
      if (firstLoadList && Array.isArray(defaultExpandedKeys) && defaultExpandedKeys.length > 0) {
        setExpandedKeys(defaultExpandedKeys);
        // 延时滑动到 defaultExpandedKeys，不然不会加载 defaultExpandedKeys，不然不会加载版本
        setTimeout(() => {
          treeRef.current?.scrollTo({ key: defaultExpandedKeys[0], align: 'bottom' });
        }, 100);
      }
    };
    restoreLastExpand();
  }, [firstLoadList, defaultExpandedKeys]);

  // 获取数据集\模型\镜像版本列表
  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 操作，需要延时以便 TreeData 更新完
      setTimeout(() => {
        restoreLastCheck(parentId, res);
      }, 100);
    } 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);
      setVersionDesc(res.versionDesc);
    } else {
      setVersionPath('');
      setFiles([]);
      setVersionDesc(undefined);
    }
  };

  // 展开时，动态加载 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([]);
      setVersionDesc(undefined);
    }
  };

  // 恢复上一次的 Check 操作
  // 判断是否有 defaultCheckedKeys，如果有，设置 checkedKeys，并且调用获取文件列表接口
  // fisrtLoadVersions 标志位
  const restoreLastCheck = (parentId: string, versions: TreeDataNode[]) => {
    if (!firstLoadVersions && Array.isArray(defaultCheckedKeys) && defaultCheckedKeys.length > 0) {
      const last = defaultCheckedKeys[0] as string;
      const { id } = getIdAndVersion(last);
      // 判断正在打开的 id 和 defaultCheckedKeys 的 id 是否一致
      if (id === parentId) {
        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);
      }
    }
  };

  // 提交
  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 childNode = treeNode.children.filter((v: TreeDataNode) => v.key === last)[0];
      const res =
        type === ResourceSelectorType.Mirror
          ? {
              activeTab: activeTab,
              ...childNode,
            }
          : {
              activeTab: activeTab,
              id,
              name,
              path: versionPath,
              version,
              identifier,
              owner,
            };
      onOk?.(res);
    } else {
      onOk?.(undefined);
    }
  };

  const title = `选择${config.name}`;
  const palceholder = `请输入${config.name}名称`;
  const fileLen = files.length > 0 ? `（${files.length}）` : '';
  const fileTitle =
    type === ResourceSelectorType.Mirror ? '镜像地址' : `${config.name}版本文件${fileLen}`;
  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 style={{ height: '50%' }}>
              <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>
            {versionDesc && (
              <div style={{ height: '50%' }}>
                <div
                  className={styles['model-selector__right__title']}
                >{`${config.name}版本描述`}</div>
                <div className={styles['model-selector__right__desc']}>{versionDesc}</div>
              </div>
            )}
          </div>
        </div>
      </div>
    </KFModal>
  );
}

export default ResourceSelectorModal;