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

import datasetImg from '@/assets/img/modal-select-dataset.png';
import mirrorImg from '@/assets/img/modal-select-mirror.png';
import modelImg from '@/assets/img/modal-select-model.png';
import { AvailableRange, CommonTabKeys, MirrorVersionStatus } from '@/enums';
import { ResourceData, ResourceVersionData } from '@/pages/Dataset/config';
import { MirrorVersionData } from '@/pages/Mirror/Info';
import { MirrorData } from '@/pages/Mirror/List';
import {
  getDatasetInfo,
  getDatasetList,
  getDatasetVersionList,
  getModelInfo,
  getModelList,
  getModelVersionList,
} from '@/services/dataset/index.js';
import { getMirrorListReq, getMirrorVersionListReq } from '@/services/mirror';
import type { TabsProps, TreeDataNode } from 'antd';
import { pick } from 'lodash';

export enum ResourceSelectorType {
  Model = 'Model', // 模型
  Dataset = 'Dataset', // 数据集
  Mirror = 'Mirror', //镜像
}

// 数据集、模型列表转为树形结构
const convertDatasetToTreeData = (list: ResourceData[]): TreeDataNode[] => {
  return list.map((v) => ({
    ...v,
    key: `${v.id}`,
    title: v.name,
    isLeaf: false,
    checkable: false,
  }));
};

// 镜像列表转为树形结构
const convertMirrorToTreeData = (list: MirrorData[]): TreeDataNode[] => {
  return list.map((v) => ({
    key: `${v.id}`,
    title: v.name,
    isLeaf: false,
    checkable: false,
  }));
};

// 数据集、模型版本列表转为树形结构
const convertResourceVersionToTreeData = (
  parentId: string,
  info: ResourceData,
  list: ResourceVersionData[],
): TreeDataNode[] => {
  return list.map((item: ResourceVersionData) => ({
    ...pick(info, ['id', 'name', 'owner', 'identifier', 'is_public']),
    version: item.name,
    title: item.name,
    key: `${parentId}-${item.name}`,
    isLeaf: true,
    checkable: true,
  }));
};

// 镜像版本列表转为树形结构
const convertMirrorVersionToTreeData = (
  parentId: string,
  list: MirrorVersionData[],
): TreeDataNode[] => {
  return list.map((item: MirrorVersionData) => ({
    url: item.url,
    title: item.tag_name,
    key: `${parentId}-${item.id}`,
    isLeaf: true,
    checkable: true,
  }));
};

interface SelectorTypeInfo {
  getList: (isPublic: boolean) => Promise<any>; // 获取资源列表
  getVersions: (parentKey: string, parentNode: any) => Promise<any>; // 获取资源版本列表
  getFiles: (parentKey: string, parentNode: any) => Promise<any>; // 获取资源文件列表
  readonly modalIcon: string; // modal icon
  readonly buttonIcon: string; // button icon
  readonly name: string; // 名称
  readonly tabItems: TabsProps['items']; // tab 列表
  readonly buttontTitle: string; // 按钮 title
}

// 数据集选择
export class DatasetSelector implements SelectorTypeInfo {
  readonly name = '数据集';
  readonly modalIcon = datasetImg;
  readonly buttonIcon = 'icon-xuanzeshujuji';
  readonly tabItems = [
    {
      key: CommonTabKeys.Private,
      label: '我的数据集',
    },
    {
      key: CommonTabKeys.Public,
      label: '公开数据集',
    },
  ];
  readonly buttontTitle = '选择数据集';

  async getList(isPublic: boolean) {
    const res = await getDatasetList({ is_public: isPublic, page: 0, size: 2000 });
    if (res && res.data) {
      const list = res.data.content || [];
      return convertDatasetToTreeData(list);
    } else {
      return Promise.reject('获取数据集列表失败');
    }
  }
  async getVersions(parentKey: string, parentNode: ResourceData) {
    const res = await getDatasetVersionList(pick(parentNode, ['owner', 'identifier']));
    if (res && res.data) {
      const list = res.data;
      return convertResourceVersionToTreeData(parentKey, parentNode, list);
    } else {
      return Promise.reject('获取数据集版本列表失败');
    }
  }

  async getFiles(_parentKey: string, parentNode: ResourceData & ResourceVersionData) {
    const params = pick(parentNode, ['owner', 'identifier', 'id', 'name', 'version', 'is_public']);
    const res = await getDatasetInfo(params);
    if (res && res.data) {
      const path = res.data.relative_paths || '';
      const list = res.data.dataset_version_vos || [];
      return {
        path,
        content: list,
      };
    } else {
      return Promise.reject('获取数据集文件列表失败');
    }
  }
}

// 模型选择
export class ModelSelector implements SelectorTypeInfo {
  readonly name = '模型';
  readonly modalIcon = modelImg;
  readonly buttonIcon = 'icon-xuanzemoxing';
  readonly tabItems = [
    {
      key: CommonTabKeys.Private,
      label: '我的模型',
    },
    {
      key: CommonTabKeys.Public,
      label: '公开模型',
    },
  ];
  readonly buttontTitle = '选择模型';

  async getList(isPublic: boolean) {
    const res = await getModelList({ is_public: isPublic, page: 0, size: 2000 });
    if (res && res.data) {
      const list = res.data.content || [];
      return convertDatasetToTreeData(list);
    } else {
      return Promise.reject('获取模型列表失败');
    }
  }
  async getVersions(key: string, parentNode: ResourceData) {
    const res = await getModelVersionList(pick(parentNode, ['owner', 'identifier']));
    if (res && res.data) {
      const list = res.data;
      return convertResourceVersionToTreeData(key, parentNode, list);
    } else {
      return Promise.reject('获取模型版本列表失败');
    }
  }

  async getFiles(_parentKey: string, parentNode: ResourceData & ResourceVersionData) {
    const params = pick(parentNode, ['owner', 'identifier', 'id', 'name', 'version', 'is_public']);
    const res = await getModelInfo(params);
    if (res && res.data) {
      const path = res.data.relative_paths || '';
      const list = res.data.model_version_vos || [];
      return {
        path,
        content: list,
      };
    } else {
      return Promise.reject('获取模型文件列表失败');
    }
  }
}

// 镜像选择
export class MirrorSelector implements SelectorTypeInfo {
  readonly name = '镜像';
  readonly modalIcon = mirrorImg;
  readonly buttonIcon = 'icon-xuanzejingxiang';
  readonly tabItems = [
    {
      key: CommonTabKeys.Private,
      label: '我的镜像',
    },
    {
      key: CommonTabKeys.Public,
      label: '公开镜像',
    },
  ];
  readonly buttontTitle = '选择镜像';

  async getList(isPublic: boolean) {
    const res = await getMirrorListReq({
      image_type: isPublic ? AvailableRange.Public : AvailableRange.Private,
      page: 0,
      size: 2000,
    });
    if (res && res.data) {
      const list = res.data.content || [];
      return convertMirrorToTreeData(list);
    } else {
      return Promise.reject('获取镜像列表失败');
    }
  }
  async getVersions(parentKey: string) {
    const res = await getMirrorVersionListReq({
      image_id: parentKey,
      page: 0,
      size: 2000,
      status: MirrorVersionStatus.Available,
    });
    if (res && res.data) {
      const list = res.data.content || [];
      return convertMirrorVersionToTreeData(parentKey, list);
    } else {
      return Promise.reject('获取镜像版本列表失败');
    }
  }

  async getFiles(_parentKey: string, parentNode: MirrorVersionData) {
    const { url } = parentNode;
    return {
      path: url,
      content: [
        {
          url: url,
          file_name: `${url}`,
        },
      ],
    };
  }
}

export const selectorTypeConfig: Record<ResourceSelectorType, SelectorTypeInfo> = {
  [ResourceSelectorType.Model]: new ModelSelector(),
  [ResourceSelectorType.Dataset]: new DatasetSelector(),
  [ResourceSelectorType.Mirror]: new MirrorSelector(),
};