#
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
# 提取自: verilog_tools_class.py
# 用途: 波形追踪工具 - 完整封装
#

import os
import re
from typing import Annotated
from debug_graph_analyzer import DebugGraph
from vcd_waveform_analyzer import parse_mismatch, get_tabular


def check_functionality(vvp_output: str) -> bool:
    """检查仿真输出是否有mismatch"""
    lines = vvp_output.splitlines()
    mismatches = None
    for line in lines:
        if re.match("^Mismatches:", line):
            mismatches = int(line.split()[1])
            break
        elif re.match("^Hint: Total mismatched samples is ", line):
            mismatches = int(line.split()[5])
            break

    if mismatches is None:
        return True  # 没找到mismatch信息，认为通过
    return mismatches == 0


def get_input_ports(verilog_code: str) -> list:
    """从Verilog代码中提取输入端口"""
    input_ports = []
    lines = verilog_code.splitlines()
    for line in lines:
        # 匹配 input 声明
        match = re.search(r'input\s+(?:wire|reg)?\s*(?:\[[^\]]+\])?\s*(\w+)', line)
        if match:
            input_ports.append(match.group(1))
    return input_ports


class WaveformTracer:
    """波形追踪工具类"""

    def __init__(self, verilog_file_path: str, vcd_file_path: str):
        self.verilog_file_path = verilog_file_path
        self.vcd_file_path = vcd_file_path
        self.graph_tracer = None
        self.cur_verilog_content = ""

    def _load_verilog(self) -> str:
        """加载Verilog文件内容"""
        with open(self.verilog_file_path, 'r', encoding='utf-8') as f:
            return f.read()

    def waveform_trace_tool(
        self,
        function_check_output: Annotated[str, "仿真工具的输出字符串"],
        trace_level: Annotated[int, "信号回溯层数，应大于1"]
    ) -> str:
        """
        波形追踪工具 - 核心函数

        参数:
        - function_check_output: 仿真工具的输出（包含mismatch信息）
        - trace_level: BFS回溯层数

        返回:
        - 包含信号追踪链、波形表、代码片段的字符串
        """

        # 1. 检查文件是否存在
        if not os.path.exists(self.vcd_file_path):
            return "[Error] wave.vcd is not found! Please run simulation first!"
        if not os.path.exists(self.verilog_file_path):
            return "[Error] Verilog file is not found!"

        # 2. 加载Verilog代码
        verilog_content = self._load_verilog()

        # 3. 构建AST调试图（如果代码有变化则重建）
        if self.cur_verilog_content != verilog_content:
            print("Creating new AST tree graph...")
            self.graph_tracer = DebugGraph([self.verilog_file_path])
            self.cur_verilog_content = verilog_content

        # 4. 检查是否有mismatch
        print("Get mismatched signal...")
        if check_functionality(function_check_output):
            print("No mismatched signals")
            return "[Waveform Tracer]: No mismatched signals!"

        # 5. 解析mismatch信息
        mismatch_columns, offset = parse_mismatch(test_output=function_check_output)

        # 6. BFS回溯，找到控制信号
        print("Trace graph signal...")
        traced_signals_map, signal_level_tracer = self.graph_tracer.get_k_control_signals(
            target_signals=mismatch_columns,
            k=trace_level,
            signal_only=True
        )

        # 7. 格式化信号追踪链
        traced_signal_str = "[Signal Traces] Backtrace control signal relations.\n"
        for bt in range(len(signal_level_tracer) - 1, -1, -1):
            if bt == len(signal_level_tracer) - 1:
                for signal_rel in signal_level_tracer[bt]:
                    traced_signal_str += signal_rel + "\n"

            header_space = "-" * (len(signal_level_tracer) - 1 - bt)
            for signal_rel in signal_level_tracer[bt]:
                traced_signal_str += header_space + signal_rel
                if bt == 0:
                    traced_signal_str += " (*last output port level)"
                traced_signal_str += "\n"
        traced_signal_str += "\n"

        # 8. 收集所有需要追踪的信号
        all_traced_signals = [str(k) for k in traced_signals_map.keys()]

        # 9. 添加输入端口
        input_ports = get_input_ports(verilog_content)
        for inp in input_ports:
            mismatch_columns.append(inp)
            if inp not in all_traced_signals:
                all_traced_signals.append(inp)

        # 10. 获取波形表
        print("Get table waveform...")
        waveform_table_str = get_tabular(
            method='dataframe',
            vcd_path=self.vcd_file_path,
            mismatch_columns=all_traced_signals,
            offset=offset,
            ori_mismatch_columns=mismatch_columns
        )

        # 11. 格式化波形信息
        waveform_table_str = (
            "[Signal Waveform]: <signal>_tb is the given testbench signal and can not be changed! "
            "<signal>_ref is the golden, and <signal>_dut is the generated verilog file waveform. "
            "Check the mismatched signal waveform and its traced signals. "
            "The clock cycle (clk) is 10ns and toggles every 5ns.\n"
            "'-' means unknown during simulation. "
            "If the '-' is the reason of mismatched signal, please check the reset and assignment block.\n"
            f"[Testbench Input Port Signal to Module]: {', '.join(input_ports)}\n"
            f"[Traced Signals]: {', '.join(all_traced_signals)}\n"
            f"[Table Waveform in hexadecimal format]\n{waveform_table_str}"
        )

        # 12. 读取Verilog代码
        with open(self.verilog_file_path, 'r', encoding='utf-8') as f:
            verilog_lines = f.readlines()

        logic_str = "[Verilog of DUT]:\n```verilog\n" + ''.join(verilog_lines) + "\n```\n"

        # 13. 生成提示信息
        HINT = (
            f"\n\n[Note] You can not change the [testbench input signal]: ({', '.join(input_ports)})! "
            "Modify the module implementation considering the input signals.\n"
            "[Hint] Firstly, identify the time of mismatched signals, and only focus on the mismatched signals in the waveform firstly. "
            "Then, explain the related signals and their transitions in the waveform table. "
            "Don't correct signals without mismatch in the table waveforms. "
            "If the information is not enough for correct the functional error, "
            f"try to trace more relevant signals using trace_level > {trace_level} for waveform_trace_tool."
        )

        # 14. 如果mismatch发生在开始时，添加初始化提示
        if offset <= 5:
            HINT += (
                "\n\n[Debug report]: The mismatch happened at the beginning. "
                "Check and set the correct initial value for mismatched signal.\n"
                "### Example of initialize the signal to 0 ###\n\n"
                "logic [3:0] a;\ninitial\n  a=4'b0000;\n\n"
                "### End example ###"
            )

        return logic_str + waveform_table_str + HINT


# ============ 使用示例 ============
if __name__ == '__main__':
    # 示例用法
    tracer = WaveformTracer(
        verilog_file_path="path/to/your/module.v",
        vcd_file_path="path/to/your/wave.vcd"
    )

    # 假设这是仿真输出
    sim_output = """
    Mismatches: 3
    Time 100ns: count expected=0001, got=0000
    """

    result = tracer.waveform_trace_tool(
        function_check_output=sim_output,
        trace_level=2
    )
    print(result)
