# -------------------------------------------------------------------------
#  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 importlib
try:
    import torch_npu
except ImportError:
    pass

from atk.configs.dataset_config import InputDataset
from atk.tasks.api_execute import register
from atk.tasks.api_execute.base_api import BaseApi
from atk.tasks.api_execute.aclnn_base_api import AclnnBaseApi


@register("function")
class FunctionApi(BaseApi):
    def __call__(self, input_data: InputDataset, with_output: bool = False):
        if not with_output:
            func = self.get_safe_eval(self.api_name)
            func(*input_data.args, **input_data.kwargs)
            return None
        
        if self.output is None:
            func = self.get_safe_eval(self.api_name)
            output = func(*input_data.args, **input_data.kwargs)
        else:
            func = self.get_safe_eval(self.api_name)
            func(*input_data.args, **input_data.kwargs)
            if isinstance(self.output, int):
                output = input_data.args[self.output]
            elif isinstance(self.output, str):
                output = input_data.kwargs[self.output]
            else:
                raise ValueError(f"self.output value is error: {self.output}")
        return output

    def get_safe_eval(self, func_name_str):
        parts = func_name_str.split(".")
        current_object = importlib.import_module(parts[0])
        for i in range(1, len(parts)):
            current_object = getattr(current_object, parts[i])
        return current_object


@register("aclnn_function")
class AclnnFunctionApi(AclnnBaseApi):
    def __call__(self):
        super().__call__()

    def init_by_input_data(self, input_data: InputDataset):
        return super().init_by_input_data(input_data)

    def after_call(self, output_packages):
        return super().after_call(output_packages)


@register("function_list")
class FunctionListApi(FunctionApi):
    def __call__(self, input_data: InputDataset, with_output: bool = False):
        output = super(FunctionListApi, self).__call__(input_data, with_output=with_output)
        return [output] if output is not None else output