# -------------------------------------------------------------------------
#  This file is part of the MindStudio project.
# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
#
# MindStudio is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#
#          http://license.coscl.org.cn/MulanPSL2
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
# -------------------------------------------------------------------------

from atk.configs.case_config import CaseConfig
from atk.configs.nodes_config import NodesConfig
from atk.tasks.task_creator.base_task import BaseTasks
from atk.tasks.task_creator.aclnn_task import AclnnTask
from atk.tasks.backends import BackendsFactory


class TaskCreator:
    def __init__(
            self, args, nodes_config: NodesConfig
    ):
        self.args = args
        self.nodes_config = nodes_config
        self.task = None
        self.init_config()

    def init_config(self):
        """
            初始化配置
        """
        if BackendsFactory.is_aclnn(self.nodes_config.get_main_node().backend):
            task = AclnnTask
        else:
            task = BaseTasks
        self.task = task(self.args, self.nodes_config)

    def run_for_case(self, case_config, task_id=None, all_cases_num=None):
        """
            创建任务链
            return: chain()
        """
        task = self.task.create_task(
            case_config,
            task_id=task_id,
            all_cases_num=all_cases_num
        )
        return task

    def create_clean_callback_task(self, case_config: CaseConfig):
        """
            创建清理数据任务
        """
        if self.task is None:
            task = BaseTasks(self.args, self.nodes_config)
        else:
            task = self.task
        return task.create_clean_callback_task(case_config)

    def create_task_result(self, case_config: CaseConfig):
        if self.task is None:
            task = BaseTasks(self.args, self.nodes_config)
        else:
            task = self.task
        return task.create_task_result(case_config)