# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
# All rights reserved.
#
# Licensed under the Apache License, Version 2.0  (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from msprobe.core.overflow_check.level import OverflowLevel
from msprobe.visualization.graph.node_op import NodeOp
from msprobe.visualization.utils import Suggestions, GraphConst
from msprobe.visualization.builder.msprobe_adapter import format_node_data, compare_data


class BaseNode:
    def __init__(self, node_op, node_id, up_node=None):
        self.op = node_op
        self.id = node_id
        self.data = {}
        self.output_data = {}
        self.input_data = {}
        self.upnode = None
        self.add_upnode(up_node)
        self.subnodes = []
        self.matched_node_link = []
        self.suggestions = {}
        self.stack_info = []
        self.micro_step_id = None
        self.overflow_level = None

    def __str__(self):
        info = f'id:\t{self.id}'
        return info

    def __eq__(self, other):
        """
        用来判断两个节点是否可以被匹配上，认为结构上是否一致
        """
        if not compare_data(self.input_data, other.input_data):
            return False
        if not compare_data(self.output_data, other.output_data):
            return False
        return True

    def get_suggestions(self):
        """
        精度疑似有问题时，提供一些建议
        """
        if self.op == NodeOp.module:
            self.suggestions[GraphConst.SUGGEST_KEY] = Suggestions.Module
            self.suggestions[Suggestions.DUMP] = Suggestions.DUMP_URL
        elif self.op == NodeOp.function_api:
            self.suggestions[GraphConst.SUGGEST_KEY] = Suggestions.API
            self.suggestions[Suggestions.API_ACCURACY_CHECKER] = Suggestions.API_ACCURACY_CHECKER_URL

    def set_input_output(self, input_data, output_data):
        self.input_data = input_data
        self.output_data = output_data

    def set_overflow_level(self, level):
        if not level or not isinstance(level, OverflowLevel):
            return
        self.overflow_level = level

    def add_upnode(self, node):
        """
        绑定upnode，用于对两个节点进行上下级关联
        """
        if not node or node.id == self.id or self.upnode:
            return
        self.upnode = node
        node.subnodes.append(self)

    def add_link(self, node, ancestors):
        """
        在节点匹配成功后进行匹配数据的录入
        Args:
            node: 和self相互匹配的节点
            ancestors: 对面节点的祖先信息
        """
        self.matched_node_link = ancestors
        node.matched_node_link = ancestors

    def to_dict(self):
        """
        输出数据
        """
        result = {
            'id': self.id,
            'node_type': self.op.value,
            'output_data': format_node_data(self.output_data),
            'input_data': format_node_data(self.input_data),
            'upnode': self.upnode.id if self.upnode else 'None',
            'subnodes': [node.id for node in self.subnodes],
            'matched_node_link': self.matched_node_link,
            'suggestions': self.suggestions,
            'stack_info': self.stack_info
        }
        if self.micro_step_id is not None:
            result['micro_step_id'] = self.micro_step_id
        # 是否存在overflow，并保存结果
        if self.overflow_level and isinstance(self.overflow_level, OverflowLevel):
            if self.data is None:
                self.data = dict()
            self.data['overflow_level'] = self.overflow_level.value
        result['data'] = self.data
        return result

    def get_ancestors(self):
        """
        获取节点所有祖先的列表
        """
        ancestors = []
        current_node = self.upnode
        while current_node:
            ancestors.append(current_node.id)
            current_node = current_node.upnode
        return list(reversed(ancestors))
