#!/usr/bin/env python
# coding=utf-8
"""
Function:
AclOpGenerator class. This class mainly implements AscendC op call of kernel function test code generation.
# -------------------------------------------------------------------------
# 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 re
from msopst.st.interface.const_manager import ConstManager
from msopst.st.interface import utils
from msopst.template.ascendc_snippet import AscendcCodeTemplate
from msopst.st.interface.global_config_parser import GlobalConfig as GC


def _append_content_to_file(content, file_path):
    utils.print_step_log(
        "[%s] Generate test code of calling of kernel function for AscendC operator." % (os.path.basename(__file__)))
    try:
        with os.fdopen(os.open(file_path, ConstManager.WRITE_FLAGS,
                               ConstManager.WRITE_MODES), 'a+') as file_object:
            file_object.write(content)
    except OSError as err:
        utils.print_error_log("Unable to write file(%s): %s." % (file_path, str(err)))
        raise utils.OpTestGenException(ConstManager.OP_TEST_GEN_WRITE_FILE_ERROR) from err
    utils.print_info_log("Content appended to %s successfully." % file_path)


class GenAscendCOpTestCode:
    """
    Class for generating  AscendC op call of kernel function test code.
    """
    WHITE_LISTS = GC.instance().white_lists

    def __init__(self, case_list, kernel_file, output_path, machine_type):
        self.testcase_list = case_list
        self.kernel_file = kernel_file
        self.op_info = {}
        self.output_path = utils.check_output_path(
            output_path, case_list, machine_type)

    @staticmethod
    def _parser_content_by_regex(regex, content, message):
        parser_info_list = re.findall(r"%s" % regex, content)
        if not parser_info_list:
            utils.print_error_log(message)
            raise utils.OpTestGenException(ConstManager.INVALID_GEN_TEST_ASCENDC_OP_ERROR)
        return parser_info_list[0]

    @staticmethod
    def _replace_cmake_content(op_file: str, old_str: str, new_str: str) -> any:
        utils.check_path_valid(op_file, isdir=False)
        file_data = ""
        with open(op_file, 'r+', encoding="utf-8") as fout:
            all_the_lines = fout.readlines()
            fout.seek(0)
            fout.truncate()
            for line in all_the_lines:
                if old_str in line:
                    line = line.replace(old_str, new_str)
                file_data += line
        with os.fdopen(os.open(op_file, ConstManager.WRITE_FLAGS, ConstManager.WRITE_MODES), 'w') as new_fout:
            new_fout.write(file_data)

    def generate(self):
        """
        Function Description:
        generate op c++ files containing call of kernel information of AscendC kernel
        :return:
        """
        self._copy_template_to_target()
        self._parse_input_info()
        self._parse_kernel_func()
        self._write_main_cpp()
        utils.print_info_log("AscendC operator test code files for kernel implement have been successfully generated.")
        utils.print_info_log("If you want to execute kernel function in Ascend aihost or cpu, "
                             "please execute commands: cd {out} && "
                             "bash run.sh [KERNEL_NAME]({kernel_name}) [SOC_VERSION](ascend910/ascend610/ascend310p) "
                             "[CORE_TYPE](AiCore/VectorCore) [RUN_MODE](cpu/npu). "
                             "For example: cd {out} && bash run.sh {kernel_name} ascend910 AiCore npu".format(
                                out=self.output_path,
                                kernel_name=self.op_info.get('kernel_info').get('kernel_func_name')))

    def _parse_input_info(self):
        input_info_dict = self._get_input_desc(self.testcase_list[0], 'input_desc', 'input')
        output_info_dict = self._get_input_desc(self.testcase_list[0], 'output_desc', 'output')
        self.op_info.update({
            'input_desc': input_info_dict,
            'output_desc': output_info_dict
        })

    def _get_input_desc(self, testcase_struct, op_key, input_type):
        op_info_dict = {}
        for index, desc_dict in enumerate(testcase_struct[op_key]):
            input_name = desc_dict.get('name')
            if not input_name:
                utils.print_error_log("There is no name filed in json.")
                raise utils.OpTestGenException(ConstManager.INVALID_GEN_TEST_ASCENDC_OP_ERROR)
            input_shape = desc_dict.get('shape')
            shape_size = ' * '.join(map(str, input_shape))
            if input_name not in op_info_dict:
                op_info_dict[input_name] = {}
                for key in ConstManager.INPUT_INFO_KEY:
                    if key not in op_info_dict.get(input_name):
                        op_info_dict[input_name][key] = {}
            op_info_dict[input_name]['shape'] = shape_size
            op_info_dict[input_name]['type'] = self.WHITE_LISTS.ascendc_cxx_type_map.get(desc_dict.get('type'))
            file_path = os.path.join(self.output_path, 'data',
                                     testcase_struct.get('case_name') + '_' + input_type + '_' + str(index) + '.bin')
            op_info_dict[input_name]['value'] = file_path
        return op_info_dict

    def _copy_template_to_target(self):
        template_path = os.path.realpath(os.path.split(os.path.realpath(__file__))[0]
                                         + ConstManager.ASCENDC_RELATIVE_TEMPLATE_PATH)
        utils.copy_template(template_path, self.output_path, True)
        self._replace_cmake_content(os.path.join(self.output_path, 'CMakeLists.txt'), '{replace}', self.kernel_file)

    def _parse_kernel_func(self):
        # 1.read {op_name_ascendc}.cpp as a content.
        kernel_content = utils.read_file(self.kernel_file)
        if not kernel_content:
            utils.print_error_log("The kernel file is empty.")
            raise utils.OpTestGenException(ConstManager.INVALID_OP_HOST_ERROR)
        # 2. parse kernel function.
        kernel_func_info = self._parser_content_by_regex(
            'extern "C"(.+?)\n', kernel_content, "There can not parse operator description in kernel file.")
        kernel_func = self._parser_content_by_regex('void(.+?)\(', kernel_func_info,
                                                    "There is no kernel function in kernel file.")
        # 3 .find keyword '__CCE_KT_TEST__'. parse call of kernel function.
        call_kernel_info = kernel_content[kernel_content.find('__CCE_KT_TEST__'):]
        if not call_kernel_info.startswith('__CCE_KT_TEST__'):
            utils.print_error_log("There is no keyword '__CCE_KT_TEST__' in kernel file.")
            raise utils.OpTestGenException(ConstManager.INVALID_OP_HOST_ERROR)
        # 4.format content in kernel implement.
        regex = r'//.*|/\*(\s|.)*?\*/'
        content_string = re.sub(regex, '', call_kernel_info).replace('\n', ConstManager.EMPTY)
        # 5.parse call of kernel function.
        call_kernel_func = self._parser_content_by_regex('__CCE_KT_TEST__(.+?)\{', content_string,
                                                         "There can not parse operator name in kernel file.")
        call_kernel_name = self._parser_content_by_regex('void (.+?)\(', call_kernel_func,
                                                         "There is no calling of kernel function in kernel file.")
        if 'kernel_info' not in self.op_info:
            self.op_info['kernel_info'] = {}
            for key in ConstManager.KERNEL_KEY:
                if key not in self.op_info.get('kernel_info'):
                    self.op_info['kernel_info'][key] = {}
        self.op_info['kernel_info']['call_kernel_func'] = call_kernel_func
        self.op_info['kernel_info']['call_kernel_name'] = call_kernel_name.strip()
        pattern = r'stream\s*,\s*(.*?)\s*\)'
        input_info_list = re.findall(pattern, call_kernel_func.strip(), re.DOTALL)
        data_type_dict = {}
        if input_info_list:
            value_list = list(map(lambda x: x.strip().split(' '), input_info_list[0].split(',')))
            for sublist in value_list:
                if len(sublist) >= 2:
                    # set param and type, like: uint8_t* x
                    data_type_dict[sublist[1]] = sublist[0]
        # 5.parse operator description.
        self.op_info['kernel_info']['kernel_func_name'] = kernel_func.strip()
        self.op_info['kernel_info']['kernel_func'] = kernel_func_info
        self.op_info['kernel_info']['input_name'] = list(data_type_dict.keys())

    def _check_len_for_json_vs_kernel(self, index):
        if index >= len(self.op_info['kernel_info']['input_name']):
            utils.print_error_log('The input information is inconsistent '
                                  'with kernel function parameters.')
            raise utils.OpTestGenException(ConstManager.INVALID_GEN_TEST_ASCENDC_OP_ERROR)

    def _get_input_malloc(self, op_key, input_index=0):
        input_malloc_desc_list = []
        # cpu input malloc
        cpu_input_malloc_list = []
        index = 0
        for index, input_desc in enumerate(self.op_info.get(op_key).values()):
            self._check_len_for_json_vs_kernel(index)
            if op_key == 'output_desc':
                index = index + input_index + 1
                self._check_len_for_json_vs_kernel(index)
            input_name = self.op_info['kernel_info']['input_name'][index]
            if not input_desc.get('shape'):
                # when shape is [], malloc size by sizeof(type).
                input_malloc_desc_list.append(
                    AscendcCodeTemplate.MALLOC_SIZE_WITH_NO_SHAPE.format(
                        input=input_name,
                        dtype=input_desc.get('type')))
            else:
                input_malloc_desc_list.append(AscendcCodeTemplate.MALLOC_SIZE.format(
                    input=input_name,
                    shape=input_desc.get('shape'),
                    dtype=input_desc.get('type')))
            # cpu input malloc
            cpu_input_malloc_list.append(AscendcCodeTemplate.GET_INPUT_MALLOC.format(
                input=input_name))
        input_malloc_desc = ConstManager.NEXT_LINE.join(input_malloc_desc_list)
        # get cpu input malloc
        get_cpu_input_malloc = ConstManager.NEXT_LINE.join(cpu_input_malloc_list)
        return input_malloc_desc, get_cpu_input_malloc, index

    def _get_input_info(self, op_key, input_index=0):
        save_input_file_list = []
        npu_input_malloc_list = []
        npu_input_param_list = []
        index = 0
        for index, input_desc in enumerate(self.op_info.get(op_key).values()):
            if op_key == 'output_desc':
                index = index + input_index + 1
            input_name = self.op_info['kernel_info']['input_name'][index]
            # read input data according to bin file to file size malloc.
            if op_key == 'input_desc':
                save_input_file_list.append(AscendcCodeTemplate.READ_INPUT_FILE.format(
                    input=input_name,
                    filepath=input_desc.get('value')
                ))
                # npu input malloc
                npu_input_malloc_list.append(AscendcCodeTemplate.ACL_MALLOC_INPUT.format(
                    input=input_name,
                    filepath=input_desc.get('value'))
                )
            else:
                # write output data to bin file.
                save_input_file_list.append(AscendcCodeTemplate.WRITE_OUTPUT_FILE.format(
                    output=input_name,
                    filepath=input_desc.get('value')
                ))
                npu_input_malloc_list.append(AscendcCodeTemplate.ACL_MALLOC_OUTPUT.format(output=input_name))
            npu_input_param_list.append('{input}Device'.format(input=input_name))
        # read input file
        read_input_file = ConstManager.NEXT_LINE.join(save_input_file_list)
        # npu input malloc
        npu_input_malloc = ConstManager.NEXT_LINE.join(npu_input_malloc_list)
        return read_input_file, npu_input_malloc, npu_input_param_list, index

    def _copy_out_malloc(self, input_index):
        copy_out_malloc_list = []
        free_output_list = []
        for index, output_desc in enumerate(self.op_info.get('output_desc').values()):
            output_name = self.op_info['kernel_info']['input_name'][index + input_index + 1]
            copy_out_malloc_list.append(AscendcCodeTemplate.MEMCPY_HOST_TO_DEVICE.format(output=output_name))
            free_output_list.append(AscendcCodeTemplate.FREE_OUTPUT.format(
                output=output_name,
                filepath=output_desc.get('value'))
            )
        # copy out malloc
        copy_out_malloc = ConstManager.NEXT_LINE_ASCENDC.join(copy_out_malloc_list)
        # FREE_OUTPUT
        free_output = ConstManager.NEXT_LINE_ASCENDC.join(free_output_list)
        return copy_out_malloc, free_output

    def _write_main_cpp(self):
        input_malloc_desc, cpu_input_malloc, input_index = self._get_input_malloc('input_desc')
        read_input_file, npu_input_malloc, npu_input_param_list, input_index = self._get_input_info('input_desc')
        output_malloc_desc, cpu_output_malloc, _ = self._get_input_malloc('output_desc', input_index)
        write_output_file, npu_output_malloc, npu_output_param_list, _ = self._get_input_info(
            'output_desc', input_index)
        # npu input parameter
        free_param_list = []
        for param_name in self.op_info.get('kernel_info').get('input_name'):
            free_param_list.append(AscendcCodeTemplate.FREE_PARAM.format(input=param_name))
        param_info = ', '.join(self.op_info.get('kernel_info').get('input_name'))
        free_param_info = ConstManager.NEXT_LINE.join(free_param_list)
        # cpu code template
        cpu_code_template = AscendcCodeTemplate.CPU_CODE_TEMPLATE.format(
            get_cpu_input_malloc=cpu_input_malloc,
            get_cpu_output_malloc=cpu_output_malloc,
            read_input_file=read_input_file,
            kernel_func_name=self.op_info.get('kernel_info').get('kernel_func_name'),
            param_info=param_info,
            write_output_file=write_output_file,
            free_param=free_param_info
        )
        malloc_op_desc_size = AscendcCodeTemplate.MALLOC_OP_DESC_SIZE.format(
            input_malloc_desc=input_malloc_desc,
            output_malloc_desc=output_malloc_desc,
        )
        # malloc npu size
        npu_malloc_size = AscendcCodeTemplate.MALLOC_NPU_SIZE.format(
            npu_input_malloc=npu_input_malloc,
            npu_output_malloc=npu_output_malloc
        )
        copy_output_malloc, free_output = self._copy_out_malloc(input_index)
        # format template
        npu_input_param_list.extend(npu_output_param_list)
        npu_input_param = ', '.join(npu_input_param_list)
        main_content = AscendcCodeTemplate.MAIN_CPP_CONTENT.format(
            kernel_do_func=self.op_info.get('kernel_info').get('call_kernel_func'),
            kernel_func=self.op_info.get('kernel_info').get('kernel_func'),
            malloc_op_desc_size=malloc_op_desc_size,
            block_dim=1,
            cpu_code_template=cpu_code_template,
            device_id=0,
            npu_malloc_size=npu_malloc_size,
            kernel_do_func_name=self.op_info.get('kernel_info').get('call_kernel_name'),
            npu_input_param=npu_input_param,
            copy_out_malloc=copy_output_malloc,
            free_output=free_output
        )
        ascendc_main_cpp_path = self.output_path + ConstManager.ASCENDC_MAIN_CPP_RELATIVE_PATH
        _append_content_to_file(main_content, ascendc_main_cpp_path)