# -------------------------------------------------------------------------
#  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 glob
import os
import re
import json

import torch
try:
    import torch_npu
except ImportError:
    pass

from atk.common.utils import get_output_path_by_case, get_output_data_infos
from atk.common.log import Logger
from atk.common.file_check import safe_file_open
from atk.tasks.report.csv_report import CsvReport
from atk.tasks.report.report_title.base_report_title import TitleType
from atk.tasks.executors.opp_executor import OppExecutor
from atk.tasks.executors.compare_excutor import CompareExecutor
from atk.tasks.executors.dataset_executor import DatasetExecutor
from atk.tasks.backends.backend import ATKRuntimeError
from atk.tasks.backends.lib_interface.acl_wrapper import ACLRuntimeError
from atk.configs.base_config import TaskTypesSelect
from atk.configs.standard_config import PassStandard
from atk.configs.results_config import FunctionCostTimes, PerformanceResult, TaskResult, PerformanceConfig
from atk.configs.nodetype_config import NodeType

logging = Logger().get_logger()
run_times = {}


def get_task_result(data):
    ret = None
    if isinstance(data, dict):
        return data
    if isinstance(data, list):
        for i in data:
            ret = get_task_result(i)
            if ret:
                return ret
    return ret


def params_task_results_json(result_json):
    def get_result_list(result_data):
        if isinstance(result_data, dict):
            return [result_data]
        ret_list = []
        for ret in result_data:
            ret_list.extend(get_result_list(ret))
        return ret_list

    task_results = get_result_list(result_data=result_json)
    main_task_result = None

    task_results_cls = []
    for task_result in task_results:
        tr = TaskResult(**task_result)
        if tr.is_benchmark_task:
            continue
        if tr.nodes.get_main_node().get_backend_name() == tr.get_backend_name():
            main_task_result = tr
        else:
            task_results_cls.append(tr)
    for task_result in task_results_cls:
        main_task_result.performance.update(**task_result.performance.model_dump())
    logging.debug(f"main_task_result: {main_task_result}")
    return main_task_result


def csv_save(task_result_json):
    task_result = TaskResult(**task_result_json)
    report = CsvReport(task_result)
    data = {}
    report_data = report.report_title_factory.get_report_data(TitleType.API_TITLE)
    data["report_data"] = report_data["data_en"]
    data["perf_standard"] = report.perf_standard.model_dump()
    pass_standard = PassStandard()
    data["acc_pass_standard"] = pass_standard.get_acc_pass_ratio()
    data["mem_pass_standard"] = pass_standard.get_memory_pass_ratio()
    if task_result.is_benchmark_compare() and task_result.accuracy is not None:
        data["benchmark_result"] = task_result.get_benchmark_data()
    return data


def post_process(task_result_json):
    logging.debug(f"start post process: {task_result_json}")
    task_result = params_task_results_json(task_result_json)
    executor = CompareExecutor(task_result)
    try:
        user_pattern = executor.task_result.case_config.expected_error_msg
        if user_pattern:
            # 错误信息对比机制
            ret = executor.compare_error_msg()
            logging.info(f"error messages matching, task id is {executor.task_result.case_config.id}")
            executor.write_error_msg(ret)
        elif bool(set(TaskTypesSelect.ACCURACY_TASKS) & set(task_result.task_type)):
            executor.accuracy_calc()
    except Exception as e:
        logging.exception(e)
        raise e
    else:
        task_result = executor.task_result
        performance_configs = task_result.performance.performance_configs
        node = task_result.get_node()
        backend_name = node.get_backend_name()
        if not performance_configs.get(backend_name):
            performance_configs[backend_name] = PerformanceConfig(function_times=executor.run_times)
        else:
            performance_configs[backend_name].update(function_times=executor.run_times)
        return task_result.model_dump()


def create_dataset(task_result_json):
    task_result = TaskResult(**task_result_json)
    logging.debug(f"task result: {task_result}")
    executor = DatasetExecutor(task_result)
    executor.create_dataset()
    executor.save_dataset()
    fc_times = FunctionCostTimes()
    fc_times.update(**executor.run_times)

    data = dict()
    data[task_result.get_backend_name()] = PerformanceConfig(function_times=fc_times)
    task_result.performance = PerformanceResult()
    task_result.performance.performance_configs = data
    logging.debug(f"create dataset success, result: {task_result.model_dump()}")
    return task_result.model_dump()


def run_bm_task(dataset_executor, task_result_json):
    """
    单标杆对比精度标准下执行的task
    单标杆比对：采用单一的精度标杆（CPU、NPU小算子拼接）进行比对。
    """
    logging.debug(f"begin to run accuracy opp task.")
    task_result = TaskResult(**task_result_json)
    data = dataset_executor.load_dataset()
    opp_executor = OppExecutor(task_result=task_result)
    opp_executor.init_resource()
    opp_executor.before_run(data)
    _, output_info = opp_executor.run_accuracy(save_data=False, save_output_info=True)
    return output_info


def run_benchmark_task(task_result, dataset_executor):
    """
    三方对比精度标准下执行的task， 以cpu作为真值比对
    """
    logging.debug(f"begin to run benchmark task in {task_result.benchmark_device}.")
    if task_result.bm_output_path is not None:
        logging.info(f"benchmark task is passed, result will read from {task_result.bm_output_path}.")
    else:
        if NodeType(task_result.benchmark_device) == NodeType.CPU:
            task_result.backend = "cpu"
        # 初始化真值节点的output_path
        main_node = task_result.nodes.get_main_node()
        for node in task_result.nodes.nodes:
            logging.info(node)
            if node.backend == task_result.backend:
                node.output_path = main_node.output_path
                break

        opp_executor = OppExecutor(task_result=task_result)
        data = dataset_executor.load_dataset()
        data.get_benchmark_data()
        opp_executor.init_resource()
        opp_executor.before_run(data)
        opp_executor.output_dir = opp_executor.get_benchmark_dir()
        opp_executor.run_accuracy(save_output_info=True)
        logging.debug(f"run_benchmark_task run_times: {opp_executor.run_times}.")
        logging.debug(f"benchmark task run success, result save in {opp_executor.output_dir}.")


def run_opp_case(dataset_executor, task_result):
    opp_executor = OppExecutor(task_result=task_result)
    opp_executor.init_resource()
    data = dataset_executor.load_dataset()
    output_info_list_accuracy = None
    # 单标杆对比精度标准下执行task
    if task_result.is_bm_task():
        from atk.tasks.backends import BackendsFactory

        if BackendsFactory.is_aclnn(task_result.nodes.get_main_node().backend):
            output_info_list_accuracy = run_bm_task(dataset_executor, task_result.model_dump())
            opp_executor.is_save_output_info = False
        data.get_benchmark_data(is_bm_task=True)
    opp_executor.before_run(data)
    user_pattern = opp_executor.task_result.case_config.expected_error_msg
    try:
        ret, output_info_list = opp_executor.run_opp_case()
    except (ACLRuntimeError, ATKRuntimeError) as e:
        # 用户没填期望错误信息，就跳过检查、只对主节点的错误信息进行捕获
        if user_pattern is None or \
                opp_executor.task_result.nodes.get_main_node().backend != opp_executor.task_result.backend:
            logging.error(f"case {task_result.case_config.id} run opp failed: {e}")
            logging.exception(e)
            raise ATKRuntimeError(f"case {task_result.case_config.id} run opp case failed: {e}") from e
        else:
            output_info_list = None
            import traceback
            # 获取完整的堆栈跟踪信息
            stack_trace = traceback.format_exc()
            # 把报错信息和堆栈保存起来，后面对比时用
            error_content = f"Error: {str(e)}\n\nStack Trace:\n{stack_trace}"
            opp_executor.task_result.error_message = (
                opp_executor.task_result.error_message + "\n" + error_content
                if opp_executor.task_result.error_message else error_content
            )
    logging.debug(f"run_opp_case run_times: {opp_executor.run_times}")
    if task_result.has_performance_and_memory_task():
        run_performance_task(task_result, ret, opp_executor)
    task_result.update(output_info_list=output_info_list_accuracy if output_info_list_accuracy else output_info_list)


def run_performance_task(task_result, ret, opp_executor):
    logging.debug(f"performance ret: {ret}")
    performance_configs = task_result.performance.performance_configs
    node = task_result.get_node()
    perf = performance_configs.get(node.get_backend_name())
    if not perf:
        perf = performance_configs[node.get_backend_name()] = PerformanceConfig()
    perf.update(**ret.model_dump())
    perf.update(function_times=opp_executor.run_times)


def run_opp_task(task_result_json):
    """
    针对一个node执行任务的入口函数
    """
    logging.debug(f"start run_opp_task: {task_result_json}")
    task_result = TaskResult(**task_result_json)
    dataset_executor = DatasetExecutor(task_result=task_result)
    try:
        if task_result.is_benchmark_task:
            run_benchmark_task(task_result, dataset_executor)
        else:
            run_opp_case(dataset_executor, task_result)
    except FileNotFoundError as not_found_error:
        logging.error(f"FileNotFoundError: %s", not_found_error)
        # 用户指定输入data文件夹，但是文件数量跟用户指定的不符
        if task_result.use_input_data_dir:
            # 在这里报错然后退出程序
            raise FileNotFoundError(f"You are using your own input data, "
                                    f"but the numbers of data do not match your request. "
                                    f"Either edit -s and -e parameters or delete those") from not_found_error
        raise not_found_error
    except Exception as run_exec:
        logging.error(f"case {task_result.case_config.id} run opp failed: {run_exec}")
        raise RuntimeError("run opp task failed!") from run_exec
    logging.debug(f"run opp success: {task_result.model_dump()}")
    return task_result.model_dump()


def run_aclnn_task(task_result_json, self_result_json):
    """
    针对一个node执行任务的入口函数 - aclnn版
    参数:
        task_result_json: 标杆task_result
        self_result_json: 自身task_result
    """
    aclnn_task_result = TaskResult(**self_result_json)
    if isinstance(task_result_json, dict):
        task_result_json = [task_result_json]

    for other_task_json in task_result_json:
        aclnn_task_result.update(**other_task_json)

    if aclnn_task_result.nodes.get_accuracy_load_nodes():
        if aclnn_task_result.output_info_list is None:
            _get_remote_output_info(aclnn_task_result)

    return run_opp_task(aclnn_task_result.model_dump())


def clean_output_data(task_result_json):
    task_result = TaskResult(**task_result_json)
    dataset_executor = DatasetExecutor(task_result=task_result)
    dataset_executor.clean_dataset()
    executor = CompareExecutor(task_result)
    executor.clean_output_data()
    return task_result_json


def _get_remote_output_info(task_result: TaskResult):
    from atk.common.connection import RemoteManager

    main_node = task_result.nodes.get_main_node()
    load_nodes = task_result.nodes.get_accuracy_load_nodes()
    if load_nodes:
        load_node = load_nodes[0]
        base_dir = load_node.get_output_path()
        backend = load_node.get_backend_name()
        output_path = get_output_path_by_case(base_dir, backend, task_result.case_config)
        remote = RemoteManager(task_result.case_config, main_node, load_node)
        data_list = remote.glob(output_path)
        data_list += remote.get_output_info_json(output_path)
        for data_path in data_list:
            remote.download(data_path)

        task_result.case_config.downloaded = True

        local_path = remote.get_remote_in_local_dir()
        output_info_file_path = os.path.join(local_path, "output_info.json")
        if os.path.exists(output_info_file_path):
            with safe_file_open(output_info_file_path) as output_info_file:
                task_result.output_info_list = json.load(output_info_file)
        else:
            file_list = glob.glob(os.path.join(local_path, "output_*.pt"))
            pattern = re.compile(r"output_(\d+)\.pt")
            file_numbers = [int(pattern.search(file).group(1)) for file in file_list]
            sorted_files = [file for _, file in sorted(zip(file_numbers, file_list))]
            output_info = []
            for file in sorted_files:
                data = torch.load(file, map_location="cpu", weights_only=True)
                output_data_infos, _ = get_output_data_infos(data)
                output_info.extend(output_data_infos)

            task_result.output_info_list = output_info