# -------------------------------------------------------------------------
#  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 multiprocessing
import os
import time
import signal
from types import SimpleNamespace
from typing import List

from tqdm import tqdm

from atk.tasks.task_creator.task_creator import TaskCreator
from atk.tasks.report.reports import ResultsReport
from atk.configs.nodes_config import NodesConfig
from atk.configs.results_config import TaskResult, RunStatus, ProcessStatus
from atk.configs.case_config import CaseConfig
from atk.db.models import CaseResult
from atk.tasks.celery_config import db_pool, clear_celery_db
from atk.common.log import Logger
from atk.configs.base_config import TaskType

from atk.db.dao.case_result_dao import CaseResultDAO

logging = Logger().get_logger()


class ResultProcess(multiprocessing.Process):

    def __init__(self, case_configs: List[CaseConfig], nodes_config: NodesConfig, args: SimpleNamespace):
        super().__init__()
        self.daemon = True
        self.case_configs = case_configs
        self.nodes_config = nodes_config
        self.args = args
        self.case_result_dao = CaseResultDAO(db_pool)
        self.total_number = len(self.case_configs)
        self.finished_number = multiprocessing.Value('i', 0)
        self.stuck_start_time = None
        self.timeout_limit = self.args.db_timeout * 3  # 数据表长时间无新数据自动退出

        self.result_report = ResultsReport(self.nodes_config)
        base_name = os.path.basename(self.args.case) if self.args.case else os.path.basename(self.args.fast_case)
        self.report_api_name = os.path.splitext(base_name)[0]
        self.result_report.writer.set_api_name(self.report_api_name)
        self.task_creator = TaskCreator(self.args, self.nodes_config)

    def run(self):
        try:
            progress_bar = tqdm(total=len(self.case_configs), desc="ATK任务进度")
            while self.finished_number.value < self.total_number:
                need_save_list = self.case_result_dao.query_need_save_report_finished(self.args.db_timeout)
                if need_save_list:
                    self.stuck_start_time = None
                    self.update_process(need_save_list, progress_bar)
                else:
                    self._wait_and_check_timeout(progress_bar)
            progress_bar.close()
        except KeyboardInterrupt:
            logging.error("receive KeyboardInterrupt, stop get result!")
        finally:
            logging.info("start write report ...")
            try:
                self.result_report.summary()
            except KeyboardInterrupt:
                logging.error("receive KeyboardInterrupt, stop write report!")
            except Exception as e:
                msg = f"failed write report. e: {str(e)}"
                raise RuntimeError(msg) from e

    def retry_create_timeout_task(self, case_result: CaseResult):
        create_task_status = ProcessStatus.UNDO
        try:
            case_config = CaseConfig(**case_result.case_results.get('case_config'))
            logging.info(f"start retry_create_timeout_task task success, case_op: {case_config.name}_{case_config.id}")
            self.task_creator.run_for_case(case_config)
            create_task_status = ProcessStatus.FINISHED
        except Exception as task_create_exception:
            logging.exception(task_create_exception)
            create_task_status = ProcessStatus.FAILED
        finally:
            self.case_result_dao.save_dataset(case_result.case_id, case_result.case_name,
                                              create_task_status=create_task_status,
                                              retry_times=case_result.retry_times + 1)

    def update_process(self, need_save_list, progress_bar):
        # 从队列中获取数据并处理
        logging.debug(f'start save result, data length: {len(need_save_list)}')
        for data in need_save_list:
            save_report_status = ProcessStatus.UNDO
            task_result = TaskResult(**data.case_results)
            if data.run_status in [RunStatus.UNDO, RunStatus.PROCESSING] and data.retry_times == 0:
                self.retry_create_timeout_task(data)
                continue
            try:
                if data.run_status in [RunStatus.UNDO, RunStatus.PROCESSING] and data.retry_times == 1:
                    task_result.run_status = RunStatus.TIMEOUT
                    logging.warning(f'task retry timeout again. '
                                    f'case: {task_result.case_config.name}_{task_result.case_config.id}')
                elif task_result.run_status != data.run_status:
                    logging.warning(f'update_process task_result. '
                                    f'run_status: {task_result.run_status}, data.run_status: {data.run_status}')
                    task_result.run_status = data.run_status
                self.result_report.print(task_result)
                save_report_status = ProcessStatus.FINISHED
            except Exception as exc:
                save_report_status = ProcessStatus.FAILED
                logging.error(f'update process failed. '
                              f'case: {task_result.case_config.name}_{task_result.case_config.id}')
                logging.exception(exc)
            finally:
                self.case_result_dao.save_dataset(
                    case_id=data.case_id,
                    case_name=data.case_name,
                    task_result=task_result,
                    save_report_status=save_report_status
                )
                progress_bar.update(1)
                self.finished_number.value += 1
        logging.debug(f'save result current length: {self.finished_number.value}, total length {self.total_number}')

    def wait_clear_db(self, finished_number):
        while self.finished_number.value < finished_number:
            logging.debug(f'wait clear db, current length: {self.finished_number.value}, '
                          f'need length: {finished_number}')
            time.sleep(2)
        clear_celery_db()

    def _exceed_timeout(self):
        if self.stuck_start_time is None:
            self.stuck_start_time = time.time()
        elapsed = time.time() - self.stuck_start_time
        if elapsed > self.timeout_limit:
            logging.error("Result Process stuck lasted {:.1f} seconds, exceeding time limit, stop get result!"
                          .format(elapsed))
            raise TimeoutError("Result Process exceeding time limit.")

    def _wait_and_check_timeout(self, progress_bar):
        """等待结果并判断是否超时。"""
        logging.debug(f"start wait result, current length: {self.finished_number.value}")
        time.sleep(1)
        self._exceed_timeout()
        progress_bar.update(0)