#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# ----------------------------------------------------------------------------
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
# CANN Open Software License Agreement Version 2.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# 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 FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
# ----------------------------------------------------------------------------

import json
import os
import stat
import sys


def parse_ini_files(ini_files):
    """
    parse ini files to json
    Parameters:
    ----------------
    ini_files:input file list
    return:ops_info
    ----------------
    """
    tbe_ops_info = {}
    for ini_file in ini_files:
        if not os.path.exists(ini_file):
            print("ini file {} not exists!".format(ini_file))
            continue
        parse_ini_to_obj(ini_file, tbe_ops_info)
    return tbe_ops_info


def parse_ini_to_obj(ini_file, tbe_ops_info):
    """
    parse ini file to json obj
    Parameters:
    ----------------
    ini_file:ini file path
    tbe_ops_info:ops_info
    ----------------
    """
    with open(ini_file) as ini_file:
        lines = ini_file.readlines()
        op = {}
        op_name = ""
        for line in lines:
            line = line.rstrip()
            if not line:
                continue
            if line.startswith("["):
                op_name = line[1:-1]
                op = {}
                tbe_ops_info[op_name] = op
            else:
                key1 = line[: line.index("=")].strip()
                key2 = line[line.index("=") + 1 :].strip()
                key1_0, key1_1 = key1.split(".")
                if key1_0 not in op:
                    op[key1_0] = {}
                if key1_1 in op[key1_0]:
                    raise RuntimeError(
                        "Op:" + op_name + " " + key1_0 + " " + key1_1 + " is repeated!"
                    )
                op[key1_0][key1_1] = key2


def check_op_info(tbe_ops):
    """
    Check info info
    """
    print("\n\n==============check valid for ops info start==============")
    required_op_input_info_keys = ["paramType", "name"]
    required_op_output_info_keys = ["paramType", "name"]
    param_type_valid_value = ["dynamic", "optional", "required"]
    is_valid = True
    for op_key in tbe_ops:
        op = tbe_ops[op_key]

        for op_info_key in op:
            if op_info_key.startswith("input"):
                op_input_info = op[op_info_key]
                missing_keys = []
                for required_op_input_info_key in required_op_input_info_keys:
                    if required_op_input_info_key not in op_input_info:
                        missing_keys.append(required_op_input_info_key)
                if len(missing_keys) > 0:
                    print(
                        "op: "
                        + op_key
                        + " "
                        + op_info_key
                        + " missing: "
                        + ",".join(missing_keys)
                    )
                    is_valid = False
                else:
                    if op_input_info["paramType"] not in param_type_valid_value:
                        print(
                            "op: "
                            + op_key
                            + " "
                            + op_info_key
                            + " paramType not valid, valid key:[dynamic, optional, required]"
                        )
                        is_valid = False
            if op_info_key.startswith("output"):
                op_input_info = op[op_info_key]
                missing_keys = []
                for required_op_input_info_key in required_op_output_info_keys:
                    if required_op_input_info_key not in op_input_info:
                        missing_keys.append(required_op_input_info_key)
                if len(missing_keys) > 0:
                    print(
                        "op: "
                        + op_key
                        + " "
                        + op_info_key
                        + " missing: "
                        + ",".join(missing_keys)
                    )
                    is_valid = False
                else:
                    if op_input_info["paramType"] not in param_type_valid_value:
                        print(
                            "op: "
                            + op_key
                            + " "
                            + op_info_key
                            + " paramType not valid, valid key:[fix, range, list]"
                        )
                        is_valid = False
    print("==============check valid for ops info end================\n\n")
    return is_valid


def write_json_file(tbe_ops_info, json_file_path):
    """
    Save info to json file
    Parameters:
    ----------------
    tbe_ops_info: ops_info
    json_file_path: json file path
    ----------------
    """
    json_file_real_path = os.path.realpath(json_file_path)
    with open(json_file_real_path, "w") as f:
        # Only the owner and group have rights
        os.chmod(
            json_file_real_path,
            stat.S_IWGRP + stat.S_IWUSR + stat.S_IRGRP + stat.S_IRUSR,
        )
        json.dump(tbe_ops_info, f, sort_keys=True, indent=4, separators=(",", ":"))
    print("Compile op info cfg successfully.")


def parse_ini_to_json(ini_file_paths, outfile_path):
    """
    parse ini files to json file
    Parameters:
    ----------------
    ini_file_paths: list of ini file path
    outfile_path: output file path
    ----------------
    """
    tbe_ops_info = parse_ini_files(ini_file_paths)
    if not check_op_info(tbe_ops_info):
        print("Compile op info cfg failed.")
        return False
    write_json_file(tbe_ops_info, outfile_path)
    return True


if __name__ == "__main__":
    args = sys.argv

    output_file_path = "tbe_ops_info.json"
    ini_file_path_list = []

    for arg in args:
        if arg.endswith("ini"):
            ini_file_path_list.append(arg)
            output_file_path = arg.replace(".ini", ".json")
        if arg.endswith("json"):
            output_file_path = arg

    if len(ini_file_path_list) == 0:
        ini_file_path_list.append("tbe_ops_info.ini")

    if not parse_ini_to_json(ini_file_path_list, output_file_path):
        sys.exit(1)
    sys.exit(0)