# Copyright (c) Huawei Technologies Co., Ltd. 2026-2026. All rights reserved.
# MindIE 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 concurrent.futures
import logging
import threading
import time
from abc import ABC, abstractmethod
from typing import Dict, Callable, Optional, Any, List

from schedule import Scheduler


class ThreadPoolManager:
    """线程池管理器 - 支持关闭与重启"""

    def __init__(self):
        self._pools: Dict[str, concurrent.futures.ThreadPoolExecutor] = {}
        # 保存配置项便于重启
        self._pool_configs: Dict[str, dict] = {}
        self._lock = threading.RLock()
        self.logger = logging.getLogger(f"ThreadPoolManager")

    def create_pool_config(self, name: str, max_workers: int = None):
        with self._lock:
            self._pool_configs[name] = {"max_workers": max_workers}

    def create_pool(self, name: str) -> bool:
        """创建线程池"""
        with self._lock:
            if name in self._pools:
                self.logger.info(f"Pool {name} already exists")
                return True

            if name not in self._pool_configs:
                self.logger.info(f"poolConfig {name} does not exist, create err")
                return False

            try:
                pool_config = self._pool_configs.get(name)
                max_workers = pool_config.get("max_workers")
                executor = concurrent.futures.ThreadPoolExecutor(
                    max_workers=max_workers, thread_name_prefix=name
                )
                self._pools[name] = executor
                self.logger.info(f"Pool {name} created, max workers: {max_workers}")
                return True
            except Exception as e:
                self.logger.info(f"Pool {name} create failed, {e}")
                return False

    def submit(
        self, pool_name: str, fn: Callable, *args, **kwargs
    ) -> Optional[concurrent.futures.Future]:
        """提交任务到线程池"""
        with self._lock:
            if pool_name not in self._pools:
                self.logger.info(f"poolConfig {pool_name} does not exist, create err")
                return None

            try:
                return self._pools[pool_name].submit(fn, *args, **kwargs)
            except Exception as e:
                self.logger.info(f"Pool {pool_name} submit failed, {e}")
                return None

    def shutdown_pool(self, pool_name: str) -> bool:
        """关闭线程池，暴露配置方便重启"""
        with self._lock:
            if pool_name not in self._pools:
                self.logger.info(f"poolConfig {pool_name} does not exist, shutdown err")
                return False

            try:
                self._pools[pool_name].shutdown(wait=False, cancel_futures=True)
                del self._pools[pool_name]
                self.logger.info(f"Pool {pool_name} shutdown")
                return True
            except Exception as e:
                self.logger.info(f"Pool {pool_name} shutdown failed, {e}")
                return False

    def shutdown_all(self):
        """关闭所有线程池"""
        with self._lock:
            pools_to_shutdown = list(self._pools.keys())
        
        for pool in pools_to_shutdown:
            self.shutdown_pool(pool)
        
        self.logger.info("shutdown all pools")

    def restore_pool(self, pool_name: str) -> bool:
        """恢复线程池"""
        with self._lock:
            if pool_name not in self._pool_configs:
                self.logger.info(f"poolConfig {pool_name} does not exist, create err")
                return False

            # 如果池还存在直接返回
            if pool_name in self._pools:
                return True

            return self.create_pool(pool_name)

    def reset_all(self) -> bool:
        """重启所有线程池"""
        with self._lock:
            for pool_name in list(self._pool_configs.keys()):
                self.restore_pool(pool_name)

    def has_pool_config(self, pool_name: str) -> bool:
        return pool_name in self._pool_configs

    def is_pool_active(self, pool_name: str) -> bool:
        return pool_name in self._pools

    def remove(self, pool_name: str) -> bool:
        """
            彻底删除 pool 以及 pool_config 后续无法恢复
        :param pool_name:
        :return:
        """
        with self._lock:
            if pool_name not in self._pool_configs:
                self.logger.info(f"poolConfig {pool_name} does not exist. remove err")
                return False

            try:
                self.shutdown_pool(pool_name)
            except Exception as e:
                self.logger.info(f"shutdown pool {pool_name} failed, {e}")
                return False
            del self._pool_configs[pool_name]
            return True


class Task(ABC):

    @abstractmethod
    def run(self):
        pass

    def get_name(self) -> str:
        return self.__class__.__name__


class IntervalTask:

    def __init__(
        self, name: str, job: Task, interval: float, pool_name: str = "default"
    ):
        self.name = name
        self.interval = interval
        self.pool_name = pool_name

        # 正在执行内容
        self._task_obj: Task = job
        self._job = None
        self.logger = logging.getLogger(f"IntervalTask.{name}")
        self._is_registered = False

    def _run_task(self, thread_pool_manager: ThreadPoolManager):
        try:
            thread_pool_manager.submit(self.pool_name, self._task_obj.run)
        except Exception as e:
            self.logger.error(
                f"submit task {self._task_obj.get_name()} to pool {self.pool_name} error: {e}"
            )

    def register(self, thread_pool_manager: ThreadPoolManager, _schedule: Scheduler):
        """
            注册定时任务， 将任务通过线程池进行调用
        :param thread_pool_manager: 线程池管理器
        :param _schedule:  定时任务调度器
        :return:
        """
        if self._is_registered:
            self.logger.info(f"already registered task {self._task_obj.get_name()}")
            return True
        if not self._task_obj:
            self.logger.info(f"task {self.name} is None")
            return False
        self._job = _schedule.every(self.interval).seconds.do(
            self._run_task, thread_pool_manager
        )
        self.logger.info(
            f"scheduled task {self._task_obj.get_name()} registered successfully"
        )
        self._is_registered = True
        return True

    def unregister(self, _schedule: Scheduler):
        """停止任务"""
        if not self._is_registered:
            return

        _schedule.cancel_job(self._job)
        self._is_registered = False
        self.logger.info(f"unregistered task {self._task_obj.get_name()}")

    def is_running(self):
        return self._is_registered

    def get_info(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "interval": self.interval,
            "running": self.is_running(),
            "pool_name": self.pool_name,
        }


class OlcScheduleManager:
    """定时任务管理器 - 支持重启关闭"""

    _instance_lock = threading.RLock()
    _instance = None

    @classmethod
    def get_instance(cls) -> "OlcScheduleManager":
        if cls._instance is None:
            with cls._instance_lock:
                if cls._instance is None:
                    cls._instance = cls()
        return cls._instance

    def __init__(self):
        self._lock = threading.RLock()
        self._tasks: Dict[str, IntervalTask] = {}
        self._pool_manager: ThreadPoolManager = ThreadPoolManager()
        self._scheduler: Scheduler = Scheduler()
        self._running = True
        self.logger = logging.getLogger("ScheduleManager")
        self._worker_thread = threading.Thread(
            name="olc_run_loop", target=self._run_loop, daemon=True
        )
        self._worker_thread.start()

    def _run_loop(self):
        """调度器运行循环"""
        self.logger.debug("_run_loop start")
        while self._running:
            try:
                self._scheduler.run_pending()
            except Exception as e:
                self.logger.error(f"run loop error: {e}")
            finally:
                time.sleep(0.5)

    def create_task(
        self,
        name: str,
        task: Task,
        interval: float,
        pool_name: str = "default",
        max_worker: int = 1,
    ) -> bool:
        """
            创建任务，并放在某个线程池中执行，支持关闭，恢复
        :param name: 任务名
        :param task: 任务
        :param interval: 时间间隔
        :param pool_name: 线程池名称
        :param max_worker: 线程池工作数目
        :return:
        """
        if name in self._tasks:
            self.logger.warning(f"task {name} already exists")
            return False

        with self._lock:
            if not self._pool_manager.has_pool_config(pool_name):
                self._pool_manager.create_pool_config(pool_name, max_worker)
                self.logger.info(f"create pool config {pool_name} success")
            self._tasks[name] = IntervalTask(name, task, interval, pool_name)
            self.logger.info(f"create task {name} success")
            return True

    def start_task(self, name: str) -> bool:
        """
         启动任务
        :param name: 任务名称
        :return:
        """
        with self._lock:
            if name not in self._tasks:
                self.logger.warning(f"task {name} does not exist")
                return False
            task = self._tasks.get(name)
            if not self._pool_manager.is_pool_active(task.pool_name):
                self._pool_manager.create_pool(task.pool_name)
                self.logger.info(f"start pool {task.pool_name} success")
            self._tasks[name].register(self._pool_manager, self._scheduler)
            return True

    def remove_task(self, name: str) -> bool:
        """
            移除任务. 后续不能恢复
        :param name: 任务名称
        :return:
        """
        with self._lock:
            if name not in self._tasks:
                self.logger.warning(f"task {name} does not exist")
                return False
            task = self._tasks.get(name)
            task.unregister(self._scheduler)
            self._tasks.pop(name)
            self.logger.info(f"remove task {name} success")
            return True

    def stop_task(self, name: str) -> bool:
        """
            停止任务
        :param name: 任务名称
        :return:
        """
        with self._lock:
            if name not in self._tasks:
                self.logger.warning(f"task {name} does not exist")
                return False
            task = self._tasks.get(name)
            task.unregister(self._scheduler)
            self.logger.info(f"stop task {name} success")
            return True

    def start_all_tasks(self) -> Dict[str, bool]:
        """
            启动所有任务
        :return:
        """
        result = {}
        with self._lock:
            # 确保线程池已重启
            self._pool_manager.reset_all()

            for name, task in self._tasks.items():
                result[name] = task.register(self._pool_manager, self._scheduler)

            if not self._running:
                self._running = True
                self._worker_thread = threading.Thread(
                    name="olc_run_loop", target=self._run_loop, daemon=True
                )
                self._worker_thread.start()
            elif not self._worker_thread.is_alive():
                # 线程已死亡， 但是_running = True 需要重启线程
                self._worker_thread = threading.Thread(
                    name="olc_run_loop", target=self._run_loop, daemon=True
                )
                self._worker_thread.start()

        self.logger.info(f"start all tasks success, result: {result}")
        return result

    def stop_all_tasks(self):
        """
            停止所有任务, 并销毁线程池
        :return:
        """
        self._running = False

        for name, task in self._tasks.items():
            task.unregister(self._scheduler)

        # 关闭所有线程池
        self._pool_manager.shutdown_all()

        # 等待工作线程结束
        if hasattr(self, "_worker_thread") and self._worker_thread.is_alive():
            self._worker_thread.join(timeout=2)

        self.logger.info(f"stop all tasks success, result: {self._tasks}")

    def get_task_info(self, name: str) -> Optional[Dict[str, Any]]:
        """
            获取任务信息
        :param name: 任务名称
        :return:
        """
        with self._lock:
            if name not in self._tasks:
                self.logger.warning(f"task {name} does not exist")
                return None
        return self._tasks.get(name).get_info()

    def list_tasks(self) -> List[str]:
        """
            列出所有任务
        :return:
        """
        with self._lock:
            return list(self._tasks.keys())

    def is_running(self) -> bool:
        """
            是否正在运行
        :return:
        """
        with self._lock:
            return self._running

    def close(self):
        try:
            self._scheduler.clear()
            self.logger.info("scheduler clear success")
        except Exception as e:
            self.logger.error(f"close scheduler failed, {e}")
        finally:
            self.stop_all_tasks()
            self.logger.info("close task schedule success")