d-embodied-front/src/views/data-generation/index.tsx

import ColLayout from "@/components/ColLayout";
import { Button, Flex, message, Pagination, Popconfirm, Table } from "antd";
import { createStyles } from "antd-style";
import ButtonGroup from "antd/es/button/button-group";
import Search from "antd/es/input/Search";
import React, { useEffect, useRef, useState } from "react";
import CreateTaskModal from "./components/CreateTaskModal";
import { datasetApi } from "@/api/dataset";
import useToken from "antd/es/theme/useToken";
import { dataGenerationApi } from "@/api/data-generation";
import { preciseInterval } from "@/utils";
import dayjs from "dayjs";
import { debounce } from "lodash";

// 假设的任务类型接口，如果已存在请忽略
interface TaskItem {
  name: string;
  status: string; // 假设 status 是字符串类型
  type: string;
  create_time: string;
  update_time: string;
  id: number | string;
}

interface DataGenerationPageProps {}

const DataGenerationPage: React.FC<DataGenerationPageProps> = () => {
  const { styles } = useStyles();
  const [dataSource, setDataSource] = useState<TaskItem[] | undefined>(
    undefined
  );
  const [createTaskModalVisible, setCreateTaskModalVisible] = useState(false);
  const [datasetList, setDatasetList] = useState<
    { id: number; name: string; description: string }[] | undefined
  >(undefined);
  const [page, setPage] = useState(1);
  const [pageSize, setPageSize] = useState(10);
  const [total, setTotal] = useState(0);
  const [, token] = useToken();
  const PollyRef = useRef<any>(null);
  const refreshTime = 5000;
  const [taskTypeOptions, setTaskTypeOptions] = useState<
    { value: string; label: string }[]
  >([]);
  const [searchKeyword, setSearchKeyword] = useState<string>("");

  const HeaderSlot = (
    <div>
      <h1>数据生成</h1>
    </div>
  );

  const FooterSlot = (
    <div className={styles.pagination}>
      <Pagination
        current={page}
        pageSize={pageSize}
        total={total}
        onChange={(page) => {
          setPage(page);
        }}
        onShowSizeChange={(_, pageSize) => {
          setPageSize(pageSize);
        }}
        showSizeChanger
        showQuickJumper
        showTotal={(total) => `共 ${total} 条`}
      />
    </div>
  );

  const columns = [
    {
      title: "任务名称",
      dataIndex: "task_name",
      key: "name",
    },
    {
      title: "任务状态",
      dataIndex: "status",
      key: "status",
    },
    {
      title: "任务类型",
      dataIndex: "task_type",
      key: "type",
    },
    {
      title: "任务创建时间",
      dataIndex: "created_at",
      key: "created_at",
      render: (text: string) => dayjs(text).format("YYYY-MM-DD HH:mm:ss"),
    },
    {
      title: "任务更新时间",
      dataIndex: "updated_at",
      key: "updated_at",
      render: (text: string) => dayjs(text).format("YYYY-MM-DD HH:mm:ss"),
    },
    {
      title: "操作",
      key: "action",
      render: (_: any, record: TaskItem) => (
        <span>
          {record.status === "running" ? ( // 假设 "running" 表示执行中
            <Flex gap={16}>
              <Popconfirm
                title="确定停止吗？"
                onConfirm={() => onOperate?.("stop", record)}
              >
                <Button type="primary" danger>
                  停止
                </Button>
              </Popconfirm>
              <Button type="primary">日志</Button>
            </Flex>
          ) : (
            <Flex gap={16}>
              <Button type="primary" danger>删除</Button>
              <Button type="primary">日志</Button>
            </Flex>
          )}
        </span>
      ),
    },
  ];

  const onOperate = async (type: "stop" | "delete", record: TaskItem) => {
    try {
      if (type === "stop") {
        await dataGenerationApi.stopTask({ task_id: record.id });
      } else if (type === "delete") {
        await dataGenerationApi.deleteTask({ task_id: record.id });
      }
      setSearchKeyword('');
      setPage(1);
      getTaskList(1, pageSize, '');
      message.success("操作成功");
    } catch (error) {
      console.error(error);
    }
  };

  const getDatasetList = async () => {
    const { data } = await datasetApi.getDatasetList({
      current_page: 1,
      page_size: 100,
    });
    setDatasetList(data.items);
  };

  const getTaskList = async (
    page: number,
    pageSize: number,
    keyword?: string
  ) => {
    const { data } = await dataGenerationApi.getTaskList({
      current_page: page,
      page_size: pageSize,
      search_keyword: keyword,
    });
    setDataSource(data.items);
    setTotal(data.total);
  };

  const getOptions = async () => {
    const { data } = await dataGenerationApi.getOptions();
    setTaskTypeOptions(
      data.options.robotwin.task_types.map((item: any) => ({
        value: item,
        label: item,
      }))
    );
  };

  const clearPolling = () => {
    PollyRef.current?.cancel();
    PollyRef.current = null;
  };

  const startPolling = (page: number, pageSize: number, keyword?: string) => {
    PollyRef.current = preciseInterval(() => {
      getTaskList(page, pageSize, keyword);
    }, refreshTime);
  };

  useEffect(() => {
    clearPolling();
    getDatasetList();
    getOptions();
  }, []);

  const debounceStartPolling = debounce(
    (page: number, pageSize: number, keyword?: string) => {
      clearPolling();
      startPolling(page, pageSize, keyword);
    },
    500
  );

  useEffect(() => {
    debounceStartPolling(page, pageSize, searchKeyword);
  }, [page, pageSize, searchKeyword]);

  return (
    <>
      <ColLayout HeaderSlot={HeaderSlot} FooterSlot={FooterSlot}>
        <div className={styles.wrapper}>
          <Flex justify="space-between" align="center" className="operate">
            <Search
              placeholder="请输入任务名称模糊搜索"
              onSearch={(value) => {
                setSearchKeyword(value);
                getTaskList(page, pageSize, value);
              }}
              style={{ width: "300px" }}
            />
            <ButtonGroup>
              <Button
                type="primary"
                onClick={() => setCreateTaskModalVisible(true)}
              >
                创建任务
              </Button>
            </ButtonGroup>
          </Flex>
          <Table columns={columns} dataSource={dataSource}></Table>
        </div>
      </ColLayout>
      {createTaskModalVisible && (
        <CreateTaskModal
          visible={createTaskModalVisible}
          onCancel={() => setCreateTaskModalVisible(false)}
          onConfirm={() => setCreateTaskModalVisible(false)}
          datasetList={datasetList}
          taskTypeOptions={taskTypeOptions}
          onDatasetCreate={() => getDatasetList()}
        ></CreateTaskModal>
      )}
    </>
  );
};

export default DataGenerationPage;

const useStyles = createStyles(({ token }) => {
  return {
    wrapper: {
      height: "100%",
      width: "100%",
      overflow: "hidden",
      "& .operate": {
        marginBottom: token.margin,
      },
    },
    pagination: {
      position: "absolute",
      bottom: token.margin,
      right: token.margin,
    },
  };
});