# -------------------------------------------------------------------------
#  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
import time
import traceback

from prettytable import PrettyTable
import pandas as pd

from atk.configs.nodes_config import NodesConfig, Node
from atk.configs.case_config import CaseConfig
from atk.tasks.report.csv_report import CSVWriter
from atk.configs.results_config import (
    TaskResult,
    PerformanceConfig,
    PerformanceResult,
    PerformanceE2EConfig,
)
from atk.configs.report_config import ResultReportConfig
from atk.configs.nodetype_config import NodeType
from atk.common.log import Logger

logging = Logger().get_logger()

BACKEND_LOAD_INFO = dict(
    e2e_perf="{backend}_(.*)端到端性能",
    device_perf="{backend}_(.*)Device性能",
    memory_allocated="{backend}_(.*)分配内存",
    memory_reserved="{backend}_(.*)预留内存",
)
CPU_LOAD_INFO = dict(cpu_e2e_perf="{backend}_(.*)端到端性能")


class ResultsReport:
    def __init__(
        self,
        nodes_config: NodesConfig,
    ):
        self.results = {}
        self.nodes_config = NodesConfig(**nodes_config.model_dump())

        # 收集所有带有 bm_file 的节点
        bm_nodes = self.nodes_config.get_bm_file_nodes()

        # 检查是否有多个 bm_file
        if len(bm_nodes) > 1:
            raise ValueError(
                "Only at most one node is allowed to have the 'bm_file' attribute, please check your nodes.yaml"
            )

        # 如果有且只有一个 bm_file，则加载
        if bm_nodes:
            self._load_by_path(bm_nodes[0])

        self.functions_times = {}
        self.writer = CSVWriter(self.nodes_config)
        self.start_time = time.perf_counter()

    @staticmethod
    def get_data_by_data_row(backend, data_row):
        data = {}
        load_info = CPU_LOAD_INFO if backend == NodeType.CPU else BACKEND_LOAD_INFO
        for key, value in data_row.items():
            for data_key, zh_name in load_info.items():
                if data_key in data:
                    continue
                if re.match(zh_name.format(backend=backend), key):
                    data[data_key] = value
                    break
        return data

    def add_task_result(self, task_result: TaskResult):
        task = self.results.get(task_result.case_config.id, None)
        if not task:
            task = self.results[task_result.case_config.id] = task_result
        else:
            task.update(**task_result.model_dump())
            task.replace(task_result)
        task.nodes = self.nodes_config
        self.writer.add_and_write_case(task)
        return task

    def add_csv_all(self, task_results):
        for task in task_results:
            self.writer.add_and_write_case(task)
        self.writer.save_csv()

    def print_function_times(self, task_report: ResultReportConfig):
        titles = [f"id_{task_report.id}"]
        titles.extend(task_report.function_titles())
        table = PrettyTable(titles)
        backends = ["npu", "cpu"]
        for backend in backends:
            row = [backend]
            for title in titles[1:]:
                if title not in self.functions_times.keys():
                    self.functions_times[title] = 0
                backend_title = f"{backend}_{title}"
                if backend_title not in self.functions_times.keys():
                    self.functions_times[backend_title] = 0
                self.functions_times[backend_title] += task_report.get_function_time(
                    backend, title
                )
                row.append(self.functions_times[backend_title])
            table.add_row(row)
        table.float_format = ".4"
        logging.info(f"\n{table}")
        logging.info(f"E2E time {time.perf_counter() - self.start_time:.4f}",)

    def add_csv(self, task_result: TaskResult):
        try:
            self.writer.add_and_write_case(task_result)
        except Exception as error:
            logging.error(traceback.format_exc())
            logging.error("save csv failed")

    def print(self, task_result: TaskResult, print_log=True):
        task = self.add_task_result(task_result)
        if print_log:
            report = ResultReportConfig(task)
            logging.info(f"\n{report.get_print_table()}")

    def summary(self):
        self.writer.summary()

    def _load_by_path(self, bm_node: Node):
        path = bm_node.bm_file
        if not os.path.exists(path):
            raise FileNotFoundError(f"{path} is not exited, load benchmark file failed.")

        data_list = pd.read_excel(path, sheet_name="statistic")
        for _, data_row in data_list.iterrows():
            _id = data_row["编号"]
            case_config = CaseConfig(id=_id)
            task_ret = TaskResult(case_config=case_config, nodes=self.nodes_config)
            performance_configs = dict()
            data = ResultsReport.get_data_by_data_row(bm_node.backend, data_row)
            performance_configs[bm_node.get_backend_name()] = PerformanceConfig(
                node=bm_node,
                device_perf=data.get("device_perf"),
                e2e_config=PerformanceE2EConfig(e2e_perf=data.get("e2e_perf")),
                memory_allocated=data.get("memory_allocated"),
                memory_reserved=data.get("memory_reserved"),
            )
            task_ret.performance = PerformanceResult()
            task_ret.performance.performance_configs = performance_configs
            self.results[_id] = task_ret
        logging.debug("Init benchmark file successfully.")