# -------------------------------------------------------------------------
#  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 os
import time
import json
import shutil
from abc import ABC, abstractmethod
import torch

from atk.configs.dataset_config import InputDataset
from atk.configs.results_config import TaskResult
from atk.tasks.dataset.base_dataset import OpsDataset
from atk.tasks import record_times
from atk.common.log import Logger
from atk.common.file_check import safe_file_open
from atk.case_generator.generator.data_types import DataTypeFactory
from atk.common.utils import get_input_path_by_case, get_file_md5


logging = Logger().get_logger()


class BaseDatasetExecutor(ABC):
    def __init__(self, task_result: TaskResult):
        super().__init__()
        self.task_result = task_result

    @abstractmethod
    def create_dataset(self):
        pass

    @abstractmethod
    def save_dataset(self):
        pass

    @abstractmethod
    def clean_dataset(self):
        pass

    @abstractmethod
    def load_dataset(self):
        pass


class DatasetExecutor:
    run_times = {}

    def __init__(self, task_result: TaskResult):
        super().__init__()
        self.task_result = task_result
        self.input_dataset = InputDataset()
        self.api_type = self.task_result.case_config.api_type
        self.case_config = self.task_result.case_config
        if self.task_result.nodes:
            self.group_nodes = self.task_result.nodes.group_node_by_same_device()

    @staticmethod
    def get_input_file_path(task_result, final_input: bool = False, custom_base: str = None):
        if custom_base is not None:
            input_dir = os.path.join(custom_base, str(task_result.case_config.id))
        else:
            input_dir = task_result.get_input_final_dir() if final_input else task_result.get_input_dir()
        suffix = "final.bin" if final_input else ".bin"
        file_path = os.path.join(input_dir, f"input{suffix}")
        other_path = None
        if "tensor" in task_result.case_config.api_type:
            other_path = os.path.join(input_dir, f"tensor_input{suffix}")
        elif "method" in task_result.case_config.api_type:
            other_path = os.path.join(input_dir, f"method_input{suffix}")
        return file_path, other_path

    @staticmethod
    def _load_file_and_backward(path, inputs):
        """
        加载落盘的输入数据和输入的backward信息
        """
        if not os.path.exists(path):
            raise FileNotFoundError(f"input file {path} not found")
        data = torch.load(path, map_location="cpu", weights_only=True)
        logging.debug(f"load file success: {path}")
        args, backward_args = [], []
        kwargs, backward_kwargs = dict(), dict()
        for i, input_info in enumerate(inputs):
            if isinstance(input_info, list):
                name = input_info[0].name
                backward = input_info[0].backward
                is_tuple = isinstance(data[i], tuple)
                if is_tuple:
                    data[i] = list(data[i])
                for index, input_config in enumerate(input_info):
                    data[i][index] = (DataTypeFactory.create_customer_datatype(input_config.dtype)
                                      (input_config).get_data(data[i][index]))
                if is_tuple:
                    data[i] = tuple(data[i])
            else:
                name = input_info.name
                backward = input_info.backward
                data[i] = (DataTypeFactory.create_customer_datatype(input_info.dtype)
                           (input_info).get_data(data[i]))
            if not name:
                args.append(data[i])
                backward_args.append(backward)
            else:
                kwargs[name] = data[i]
                backward_kwargs[name] = backward
        return args, kwargs, backward_args, backward_kwargs

    @record_times(run_times)
    def create_dataset(self):
        # 如果使用用户给定的路径作为数据输入，那么跳过数据生成
        if self.task_result.use_input_data_dir:
            return
        ops = OpsDataset(
            [self.task_result.case_config.model_dump()],
            self.task_result.case_config.id,
        )
        self.input_dataset = ops.__next__()

    @record_times(run_times)
    def save_dataset(self):
        file_path, other_path = self.get_input_file_path(self.task_result)
        # 如果使用用户给定的路径作为数据输入，那么跳过数据生成，跳过数据保存
        if self.task_result.use_input_data_dir:
            return file_path, other_path
        self.input_dataset.save_input_data(file_path, other_path, self.task_result)
        return file_path, other_path

    @record_times(run_times)
    def load_dataset(self) -> InputDataset:
        # 如果 self.task_result.use_ipnut_data_dir 有内容，那么 get_input_file_path 会根据传入的input_data_dir作为输入数据的开头
        # 从而获取到用户指定的包含.bin 文件的文件夹
        file_path, other_path = self.get_input_file_path(
            self.task_result,
            custom_base=self.task_result.use_input_data_dir
        )
        if not os.path.exists(file_path) and other_path is None:
            raise FileNotFoundError(f"input file_path: {file_path} is not exists and other_path is none")
        data = InputDataset()
        if os.path.exists(file_path):
            data.args, data.kwargs, data.backward_args, data.backward_kwargs = DatasetExecutor._load_file_and_backward(
                file_path, self.case_config.inputs
            )
            logging.debug(f"load data from {file_path} success.")
        else:
            logging.warning(f"input file_path: {file_path} is not exists, please check")

        if other_path is not None:
            if os.path.exists(other_path):
                # method input不需要设置backward参数
                if "method" in self.api_type and self.case_config.method_inputs:
                    data.method_args, data.method_kwargs, _, _ = DatasetExecutor._load_file_and_backward(
                        other_path, self.case_config.method_inputs
                    )
                    logging.debug(f"load other data from {other_path} success.")
                elif "tensor" in self.api_type:
                    data.tensor_args = torch.load(other_path, map_location="cpu", weights_only=True)
                    data.backward_tensor_args = getattr(self.case_config.tensor_input, 'backward', None)
                    logging.debug(f"load other data from {other_path} success.")
                else:
                    logging.warning("not set method or tensor in api_type, please check")
            else:
                raise FileNotFoundError(f"input other_path: {other_path} is not exists")
        return data

    def clean_dataset(self):

        def _remove_by_path(file_path):
            try:
                shutil.rmtree(file_path)
                logging.debug(f"delete file {file_path} success.")
            except Exception as exc:
                logging.warning(f"Delete file {file_path} exception: {str(exc)}")

        file_path, _ = self.get_input_file_path(self.task_result)
        input_path = os.path.dirname(file_path)
        if self.task_result.save_data and "input" in self.task_result.save_data:
            logging.info(f"save input data in {input_path}")
            return
        _remove_by_path(input_path)
        logging.debug(f"clean input file success {input_path}")

    @record_times(run_times)
    def update_md5_and_wait_for_opp_task(self, timeout=1800):
        """
        remote侧计算并记录单个case的所有输入的md5值, 主节点会通过http请求获取这些值
        remote通过获取本地是否有'data_upload_completed.txt'来判断主节点的数据是否传输完成
        """
        md5_value = {}
        # 获取当前节点对应的文件目录
        backend = self.task_result.backend
        name = self.task_result.name
        input_path = self.task_result.nodes.get_node(backend=backend, name=name).get_input_path()
        dire = get_input_path_by_case(
            input_path,
            self.case_config,
            is_create=True
        )
        for root, _, files in os.walk(dire):
            for file in files:
                # 文件名为input.bin/tensor_input.bin/method_input.bin的文件进行同步
                if "input.bin" not in file:
                    continue
                local_file = os.path.join(root, file)
                local_md5 = get_file_md5(local_file)
                md5_value[file] = local_md5
        input_dir = self.task_result.get_input_dir()
        input_path = os.path.join(input_dir, f"md5.json")
        with safe_file_open(input_path, "w") as f:
            json.dump(md5_value, f)
        logging.debug("end update inputs md5 in remote node and waiting for 'data_upload_completed.txt'")
        start_time = time.time()
        while time.time() - start_time < timeout:
            # 检查文件是否存在
            file_path = os.path.join(input_dir, 'data_upload_completed.txt')
            if os.path.exists(file_path):
                logging.debug("end waiting for 'data_upload_completed.txt'")
                return True
            time.sleep(0.5)
        else:
            raise TimeoutError("Timeout waiting for the main node upload data to remote side")

    def check_in_main_node(self):
        return not self.group_nodes or self.task_result.get_node() in self.group_nodes[0]