# -------------------------------------------------------------------------
#  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 sys
import time
from pathlib import Path

from coloredlogs import StandardErrorHandler
from celery import signals
from celery.app import trace

from atk.tasks.celery_config import app
from atk.tasks.task_creator.worker_config import WorkerConfig
from atk.common.resource_utils import BindCpuManager
from atk.common.log import Logger

logging = Logger().get_logger()


class CeleryWorker(app.Worker):

    def on_before_init(self, quiet=False, **kwargs):
        self.quiet = quiet
        trace.setup_worker_optimizations(self.app, self.hostname)

        # this signal can be used to set up configuration for workers by name.
        signals.celeryd_init.send(
            sender=self.hostname, instance=self,
            conf=self.app.conf, options=kwargs,
        )


class Worker(multiprocessing.Process):

    def __init__(self, quiet=False, **kwargs):
        super().__init__()
        base_dir = str(Path(__file__).absolute().parents[0])
        if base_dir not in sys.path:
            sys.path.append(base_dir)
        app.loader.import_default_modules()
        self.name = kwargs.get('queues', 'atk')
        kwargs.pop("cpu_numbers")
        self.kwargs = kwargs

        self.worker = CeleryWorker(quiet=quiet, **kwargs)

    def run(self):
        bind_cpu_manager = BindCpuManager()
        cpu_ids_without_devices = bind_cpu_manager.cpu_ids_without_devices
        if cpu_ids_without_devices:
            bind_cpu_manager.push_without_devices()
            logging.debug(f"{self.name} success bind cpu {cpu_ids_without_devices}")
        logging.info(f'start run {self.kwargs}')
        try:
            self.worker.start()  # 启动 worker
        except BaseException as e:
            logging.debug(f"Worker {self.name} exception, error: {str(e)}, start stopping...")
        finally:
            self.worker.stop()  # 确保 worker 正确退出
            logging.debug(f"Worker {self.name} stopped.")

    def stop(self):
        self.terminate()


class WorkerManager:

    def __init__(self):
        self.workers = []

    @staticmethod
    def get_quiet_by_logging():
        for handler in logging.handlers:
            if isinstance(handler, StandardErrorHandler):
                return False if handler.level == Logger().std.DEBUG else True
        return True

    def create_worker(self, quiet: bool, worker_config: WorkerConfig):
        worker = Worker(quiet=quiet, **(worker_config.model_dump()))
        worker.start()
        time.sleep(0.1)  # 防止创建queue数据库错误
        self.workers.append(worker)

    def stop_all(self):
        logging.debug("start stop worker")
        for worker in self.workers:
            worker.stop()
        for worker in self.workers:
            while worker.is_alive():
                time.sleep(0.2)
        logging.debug("finish stop worker")