# -------------------------------------------------------------------------
#  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 typing import List
import torch

from atk.configs.dataset_config import InputDataset
from atk.configs.results_config import TaskResult
from atk.tasks.api_execute.base_api import BaseApi
from atk.tasks.backends.lib_interface.acl_wrapper import AclTensorStruct, AclTensorlistStruct, AclFormat


class AclnnBaseApi(BaseApi):
    def __init__(self, task_result: TaskResult, backend):
        super().__init__(task_result)
        self.backend = backend

    def __call__(self):
        self.backend.aclnn_x_get_workspace_size()
        self.backend.aclnn_x()

    def init_by_input_data(self, input_data: InputDataset):
        """
        初始化输入参数并整合算子输出到输入参数列表

        处理流程：
        1. 转换输入参数为aclnn所需的c++格式
        2. 收集算子输出元数据
        3. 将算子输出张量地址追加到输入参数

        核心参数：
        1. input_args：算子的入参列表，应严格符合算子的c++函数接口原型的参数顺序和类型，均应转换为ctypes式或AclTensorStruct(List[ctypes | AclTensorStruct])
        2. output_packages：算子的出参数据包列表，用于精度对比在调用算子后解析数据包以获取算子输出(List[AclTensorStruct])
        3. tensor数据包数据结构：
            class AclTensorStruct:
                tensor: AclTensor       # aclTensor的c++对象（真正传递给算子的直接参数）
                addr: int             # 内存地址（用于PyTorch转换时的指针操作）
                shape: List[int]      # 张量形状
                dtype: AclDataType    # 数据类型
        """
        input_args = []  # 算子的入参列表
        output_packages = []  # 算子的出参数据包列表
        # === 处理输入参数 ===
        # 将输入数据转换为aclnn所需的c++格式
        for i, arg in enumerate(input_data.args):
            data = self.backend.convert_input_data(arg, index=i)
            input_args.extend(data)

        for name, kwarg in input_data.kwargs.items():
            data = self.backend.convert_input_data(kwarg, name=name)
            input_args.extend(data)

        # === 处理算子输出 ===
        # 收集算子输出，并储存根据输出中的shape和dtype信息生成的AclTensorStruct数据结构
        # 输出数据结构说明：
        for index, output_data in enumerate(self.task_result.output_info_list):
            output = self.backend.convert_output_data(output_data, index)
            output_packages.extend(output)  # 保存完整AclTensorStruct结构
        # 将算子输出tensor追加到输入参数
        input_args.extend(output_packages)

        return input_args, output_packages

    def after_call(self, output_packages):
        output = []
        for output_pack in output_packages:
            if isinstance(output_pack, AclTensorStruct):
                output.append(self.acl_tensor_to_torch(output_pack))
            elif isinstance(output_pack, AclTensorlistStruct):
                output.append(self.acl_tensorlist_to_torch(output_pack))
                
        return output

    def torch_tensor_to_acl(self, tensor: torch.Tensor, fmt: AclFormat = AclFormat.ACL_FORMAT_ND) -> AclTensorStruct:
        return self.backend.torch_tensor_to_acl(tensor, fmt)

    def acl_tensor_to_torch(self, tensor_struct: AclTensorStruct) -> torch.Tensor:
        return self.backend.acl_tensor_to_torch(tensor_struct)

    def acl_tensorlist_to_torch(self, tensorlist_struct: AclTensorlistStruct) -> List[torch.Tensor]:
        return list(self.backend.acl_tensorlist_to_torch(tensorlist_struct))

    def get_format(self, input_data: InputDataset, index=None, name=None):
        """
        :param input_data: 参数列表
        :param index: 参数位置
        :param name: 参数名字
        :return:
        format at this index or name
        """
        try:
            return AclFormat.ACL_FORMAT_ND
        except ImportError as exc:
            raise ImportError("Can not import AclFormat from atk.tasks.backends.lib_interface.acl_wrapper") from exc

    def get_storage_shape(self, input_data: InputDataset, index=None, name=None):
        """
        :param input_data: 参数列表
        :param index: 参数位置
        :param name: 参数名字
        :return:
        storage shape at this index or name
        """
        return None