# -------------------------------------------------------------------------
#  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.
# -------------------------------------------------------------------------

import copy
import time
from typing import List

from kombu import uuid
from celery import chain, group
from kombu import uuid

from atk.configs.case_config import CaseConfig
from atk.configs.base_config import TASK_TIME_OUT, TaskType
from atk.configs.nodes_config import NodesConfig, Node
from atk.configs.nodetype_config import NodeType
from atk.configs.results_config import TaskResult
from atk.common.log import Logger
from atk.tasks.backends import BackendsFactory
from atk.tasks.task_creator.worker_config import (
    CpuRunWorkerConfig,
    DatasetWorkerConfig,
    CompareWorkerConfig,
    DeviceRunWorkerConfig,
    CleanOutputDataWorkerConfig,
)
from atk.tasks.celery_tasks import (
    celery_create_dataset,
    celery_run_opp_task,
    celery_post_process,
    celery_clean_output_data
)

logging = Logger().get_logger()


class Task:
    def __init__(self, task):
        self.task = task
        self._id = task.id if task.id else "id"
        self.create_time = time.time()

    @property
    def alive_duration(self):
        return time.time() - self.create_time

    def get_id(self):
        return self._id


class BaseTasks:
    def __init__(self, args, nodes_config: NodesConfig):
        self.nodes_config = nodes_config
        self.args = args
        self.group_nodes = self.nodes_config.group_node_by_same_device()
        self.execute_nodes = self.nodes_config.get_execute_nodes()
        self.all_cases_num = None
        self.main_dataset_task_id = None

    def get_device_id(self, case_config: CaseConfig, node: Node) -> int:
        """
        获取驱动id
        return: int
        """
        if not node.devices:
            raise RuntimeError(f"not have devices, please check environment or node config, args: {self.args}")
        index = case_config.id % len(node.devices)
        return node.devices[index]

    def create_task(
            self,
            case_config: CaseConfig,
            task_id: str = None,
            all_cases_num=None,
    ) -> Task:
        """
        创建任务
        return: Task(chain())
        """
        self.all_cases_num = all_cases_num
        logging.debug(f"create task, case_config: {case_config.id}, task_id:{task_id}")
        return Task(self.run(case_config, task_id=task_id))

    def run(self, case_config: CaseConfig, task_id=None):
        """
        创建整体的任务链
        return: chain()
        """
        is_one_node = len(self.nodes_config.nodes) == 1
        execute_task_id = task_id if is_one_node else None
        task_list = self._create_node_task(case_config, task_id=execute_task_id)
        task_params = self.create_task_result(case_config)
        compare = self._create_compare_task(task_params, task_id=task_id)
        ret = chain(task_list, compare)()
        return ret

    def create_task_result(self, case_config: CaseConfig):
        # args.save_data 是由 click callback 生成的字典, e.g.{'input': 'txt'}
        export_cfg = self.args.save_data if hasattr(self.args, 'save_data') else {}
        return TaskResult(
            task_type=self.nodes_config.get_main_node().task,
            case_config=case_config,
            nodes=self.nodes_config,
            save_data=list(export_cfg.keys()) if export_cfg else [],
            export_config=export_cfg,
            opp_perf=case_config.opp_perf,
            benchmark_device=self.args.bm_device,
            bm_output_path=self.args.bm_output_path,
            is_trace=self.args.trace,
            all_cases_num=self.all_cases_num,
            use_input_data_dir=self.args.input_data,
        )

    def create_dataset_task(self, task_params, node):
        dataset_task = self._create_dataset_task(task_params, node)
        if node == self.nodes_config.get_main_node():
            self.main_dataset_task_id = dataset_task.id
        return dataset_task

    def create_clean_callback_task(self, case_config: CaseConfig):
        task_params = self.create_task_result(case_config)
        device_list = self.nodes_config.get_nodes_device_list()
        clean_task = self._create_node_clean_output_task(task_params, device_list)
        return clean_task()

    def _create_node_task(self, case_config: CaseConfig = None, task_id=None):
        """
        创建生成数据和执行任务
        """
        ret = []
        task_params = self.create_task_result(case_config)
        is_benchmark_compare = task_params.is_benchmark_compare()
        is_create_benchmark_task = False
        for node_list in self.group_nodes:
            execute_tasks = []
            dataset_task = None
            for index, node in enumerate(node_list):
                if node.bm_file:
                    continue
                task_params.add_node(node)
                task_params.device_id = self.get_device_id(case_config, node)
                if index == 0:
                    dataset_task = self.create_dataset_task(task_params, node)
                    immutable = False
                else:
                    immutable = True

                task = self._create_execute_task(task_params, node, immutable=immutable, task_id=task_id)
                execute_tasks.append(task)
            end_node = node_list[-1]
            benchmark_device = task_params.benchmark_device
            should_process = self._should_process_benchmark_task(benchmark_device, end_node)

            if is_benchmark_compare and not is_create_benchmark_task and should_process:
                execute_tasks.append(self._create_benchmark_task(task_params, end_node, case_config))
                is_create_benchmark_task = True
            ret.append(chain(dataset_task, group(*execute_tasks)))
        return group(*ret)

    def _should_process_benchmark_task(self, benchmark_device, node):
        return ((node == self.nodes_config.get_main_node() and benchmark_device == NodeType.CPU) or
                (node != self.nodes_config.get_main_node() and
                 (node.backend == benchmark_device or node.host == self.nodes_config.get_main_node().host)))

    def _create_benchmark_task(self, task_params, _node, case_config):
        if _node == self.nodes_config.get_main_node() and task_params.benchmark_device == NodeType.CPU:
            node = copy.deepcopy(_node)
            node.backend = NodeType.CPU
        else:
            node = _node
        task_params.add_node(node)
        task_params.device_id = self.get_device_id(case_config, node)
        task_params.is_benchmark_task = True
        execute_task = self._create_execute_task(task_params, node, immutable=True)
        task_params.is_benchmark_task = False
        return execute_task

    def _create_node_clean_output_task(self, task_params: TaskResult, devices: List[str]):
        task_list = []
        for node_info in self.execute_nodes:
            if node_info.backend in devices and node_info.task:
                task_params.add_node(node_info)
                task_list.append(self._create_clean_output_data_task(task_params, node_info))
        return chain(task_list[0]) if len(task_list) == 1 else group(*task_list)

    def _create_dataset_task(self, task_params: TaskResult, node_info: Node):
        """
        创建生成数据任务
        主节点会为所有机器创建生成数据任务
        celery_create_dataset 用于本地生成数据
        """
        queue_name = DatasetWorkerConfig(node_info.host, node_info.port).queues
        celery_task = celery_create_dataset
        dynamic_soft_time_limit = self._set_task_timeout(task_params)
        dataset = celery_task.signature(
            args=(task_params.model_dump(),),
            queue=queue_name,
            immutable=True,
            options={"task_id": uuid(), "soft_time_limit": dynamic_soft_time_limit},
        )

        return dataset

    def _create_execute_task(self, task_params: TaskResult, node_info: Node, immutable: bool = False, task_id=None):
        """
        创建执行任务
        """
        device_type = BackendsFactory.get_device_type(task_params.backend)
        if device_type != NodeType.CPU and not (
                task_params.is_benchmark_task
                and task_params.benchmark_device == NodeType.CPU
        ):
            run_queue = DeviceRunWorkerConfig(
                node_info.host, node_info.port, device_id=task_params.device_id
            ).queues
        else:
            run_queue = CpuRunWorkerConfig(
                node_info.host, node_info.port
            ).queues
        celery_task = celery_run_opp_task
        task_id = task_id or uuid()
        dynamic_soft_time_limit = self._set_task_timeout(task_params)
        if immutable:
            execute_task = celery_task.signature(
                args=(task_params.model_dump(),),
                queue=run_queue,
                immutable=True,
                options={"task_id": task_id, "soft_time_limit": dynamic_soft_time_limit},
            )
        else:
            execute_task = celery_task.signature(
                queue=run_queue, options={"task_id": task_id, "soft_time_limit": dynamic_soft_time_limit}
            )
        return execute_task

    def _create_compare_task(self, task_params: TaskResult, task_id=None):
        """
        创建比较任务
        """
        node = self.nodes_config.get_main_node()
        compare_queue = CompareWorkerConfig(node.host, node.port).queues
        dynamic_soft_time_limit = self._set_task_timeout(task_params)
        compare_task = celery_post_process.signature(
            queue=compare_queue, options={"task_id": task_id or uuid(), "soft_time_limit": dynamic_soft_time_limit}
        )
        return compare_task

    def _create_clean_output_data_task(self, task_params: TaskResult, node_info: Node):
        """
        创建清理数据任务
        """
        queue_name = CleanOutputDataWorkerConfig(node_info.host, node_info.port).queues
        celery_task = celery_clean_output_data
        dynamic_soft_time_limit = self._set_task_timeout(task_params)
        clean_task = celery_task.signature(
            args=(task_params.model_dump(),),
            queue=queue_name,
            immutable=True,
            options={"task_id": uuid(), "soft_time_limit": dynamic_soft_time_limit},
        )
        return clean_task

    def _set_task_timeout(self, task_params: TaskResult):
        soft_time_limit = 600
        if TaskType.ACCURACY_DC in task_params.task_type:
            soft_time_limit = max(int(task_params.task_info.accuracy_dc_loop / 5), 1) * TASK_TIME_OUT
        if self.args.timeout is not None:
            soft_time_limit = self.args.timeout
        return soft_time_limit