# -------------------------------------------------------------------------
#  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 celery import chain, group

from atk.configs.case_config import CaseConfig
from atk.configs.results_config import TaskResult
from atk.tasks.celery_tasks import celery_run_aclnn_task
from atk.tasks.task_creator.base_task import BaseTasks
from atk.tasks.task_creator.worker_config import DeviceRunWorkerConfig
from atk.tasks.backends import BackendsFactory


class AclnnTask(BaseTasks):
    def run(self, case_config: CaseConfig, task_id=None):
        """
        创建整体的任务链
        按照任务链执行顺序, 任务包括:
            dataset_task_list: 数据的生成 保存 传输 每台机器该任务只需执行一次
            other_execute_task_list: 非aclnn backend的节点进行算子计算
            after_execute_tasks: aclnn backend的节点进行算子计算
            compare: 结果比对 只在主节点执行
        return: chain()
        """
        task_params = self.create_task_result(case_config)
        dataset_task_list, other_nodes = self._create_aclnn_dataset_task(task_params)
        other_execute_task_list = self._create_other_execute_tasks(
            task_params, case_config, other_nodes
        )
        after_execute_tasks = self._create_after_execute_task(task_params, case_config)
        compare = self._create_compare_task(task_params, task_id=task_id)
        if len(other_execute_task_list) == 0:
            if self.nodes_config.get_accuracy_load_nodes():
                ret = chain(dataset_task_list, after_execute_tasks, compare)()
            else:
                raise RuntimeError("not other task, please check node config")
        else:
            other_execute_task_list = (
                chain(other_execute_task_list[0])
                if len(other_execute_task_list) == 1
                else chain(chain(group(*other_execute_task_list)))
            )
            ret = chain(dataset_task_list, other_execute_task_list, after_execute_tasks, compare)()
        return ret

    def _create_aclnn_dataset_task(self, task_params):
        dataset_task_list = []
        other_nodes = []
        for node_list in self.group_nodes:
            task_params.add_node(node_list[0])
            dataset_task = self.create_dataset_task(task_params, node_list[0])

            dataset_task_list.append(dataset_task)

            for node in node_list:
                if not BackendsFactory.is_aclnn(node.backend):
                    other_nodes.append(node)
        return (
            (
                chain(dataset_task_list[0])
                if len(dataset_task_list) == 1
                else chain(chain(group(*dataset_task_list)))
            ),
            other_nodes,
        )

    def _create_other_execute_tasks(self, task_params, case_config, other_nodes):
        other_execute_task_list = []
        for node in other_nodes:
            # 如果存在非aclnn backend且bm_file不为None的节点, 那么该节点不应执行算子计算
            if node.bm_file:
                continue
            task_params.add_node(node)
            task_params.device_id = self.get_device_id(case_config, node)
            execute_task = self._create_execute_task(task_params, node, immutable=True)
            other_execute_task_list.append(execute_task)
        return other_execute_task_list

    def _create_after_execute_task(self, task_params: TaskResult, case_config: CaseConfig):
        execute_tasks = []
        execute_tasks.extend(self._create_aclnn_execute_task(task_params, case_config))
        if task_params.is_benchmark_compare():
            for node_list in self.group_nodes:
                end_node = node_list[-1]
                if self._should_process_benchmark_task(task_params.benchmark_device, end_node):
                    task = self._create_benchmark_task(task_params, end_node, case_config)
                    execute_tasks.append(task)
                    break
        return (
            chain(execute_tasks[0])
            if len(execute_tasks) == 1
            else chain(chain(group(*execute_tasks)))
        )

    def _create_aclnn_execute_task(self, task_params: TaskResult, case_config: CaseConfig):
        aclnn_execute_tasks = []
        for node_list in self.group_nodes:
            for node in node_list:
                # 如果存在aclnn backend且bm_file不为None的节点, 那么该节点不应执行算子计算
                if node.bm_file:
                    continue
                if BackendsFactory.is_aclnn(node.backend):
                    task_params.add_node(node)
                    task_params.device_id = self.get_device_id(case_config, node)
                    dynamic_soft_time_limit = self._set_task_timeout(task_params)
                    task_params.cpp_func_signature_type_path = self.args.cpp_func_signature_type_path
                    run_queue = DeviceRunWorkerConfig(node.host, node.port, task_params.device_id).queues
                    aclnn_execute_tasks.append(celery_run_aclnn_task.signature(
                        args=(task_params.model_dump(),),
                        queue=run_queue,
                        immutable=False,
                        options={"soft_time_limit": dynamic_soft_time_limit},
                    ))
        return aclnn_execute_tasks