#!D:/IDE/anaconda3/envs/Microscope_Scan python
# -*- coding: UTF-8 -*-
"""
@FileName：TaskPoolInterface.py
@Description：
@Author：XSH
@Time：2025/11/19 9:38
@Department：研发部
@Copyright：©2025 乐辰生物
"""
from PyQt5 import Qt
from AppInterfaces.GlobalInterface import GlobalInterface
from concurrent.futures import ThreadPoolExecutor, as_completed
import traceback
from abc import abstractmethod
import itertools
import multiprocessing

class TaskPoolInterface(GlobalInterface):
    Executor            : ThreadPoolExecutor    #
    Futures             : dict[object, object]  #
    WorkerSize          : int                   #


    def initExecutor   (self):
        if not self.hasExecutor   :
            self.initWorkerSize()
            self.Executor    = ThreadPoolExecutor(max_workers=self.WorkerSize)

    def initFutures    (self):
        if not self.hasFutures    :
            self.Futures     = {}

    def initWorkerSize (self):
        if not self.hasWorkerSize :
            num_cores = max(4, multiprocessing.cpu_count() // 2)
            self.WorkerSize  = num_cores

    def appendItem(self, func, *args, **kwargs):
        return self.submit(func, *args, **kwargs)


    def submit(self, func, *args, **kwargs):
        """
        提交单个任务
        :param func: 任务函数
        :param args: 函数位置参数
        :param kwargs: 函数关键字参数
        :return: Future对象（可用于获取结果或取消任务）
        """
        self.initFutures ()
        self.initExecutor()
        future = self.executor.submit(func, *args, **kwargs)
        # 可选：将任务函数与Future关联（便于后续追踪）
        # self.futures[future] = (func.__name__, args, kwargs)
        return future

    def submit_batch(self, func, tasks: list[tuple]) -> list:
        """
        批量提交任务（同一函数，不同参数）
        :param func: 任务函数
        :param tasks: 参数列表，每个元素为 tuple(args, kwargs) 或 (args,)
        :return: Future对象列表
        """
        self.initFutures()
        self.initExecutor()
        futures = []
        for task in tasks:
            args = task[0] if isinstance(task, (list, tuple)) else ()
            kwargs = task[1] if len(task) > 1 else {}
            future = self.executor.submit(func, *args, **kwargs)
            self.futures[future] = (func.__name__, args, kwargs)
            futures.append(future)
        return futures

    def wait_completion(self, futures: list, callback = None) -> list:
        """
        等待所有任务完成并获取结果，支持回调函数
        :param futures: Future对象列表
        :param callback: 回调函数，格式：def callback(result, func_name, args, kwargs)
        :return: 任务结果列表（按完成顺序）
        """

        results = []
        if not self.hasExecutor or not self.hasFutures:
            return results
        for future in as_completed(futures):
            try:
                result = future.result()  # 获取任务结果
                func_name, args, kwargs = self.futures.get(future, ("unknown", (), {}))
                results.append(result)
                # 执行回调函数（如果提供）
                if callback:
                    callback(result, func_name, args, kwargs)
            except Exception as e:
                # 捕获任务执行中的异常
                func_name, args, kwargs = self.futures.get(future, ("unknown", (), {}))
                error_msg = f"任务 {func_name} 执行失败: {str(e)}\n{traceback.format_exc()}"
                print(error_msg)
                results.append(None)  # 异常时结果置为None
        return results

    def shutdown(self, wait: bool = True) -> None:
        """
        关闭线程池
        :param wait: 是否等待所有任务完成后再关闭
        """
        if not self.hasExecutor or not self.hasFutures:
            return
        self.executor.shutdown(wait=wait)
        self.futures.clear()
        self.Executor = None

    @property
    def executor   (self) -> ThreadPoolExecutor    :
        return None if not self.hasExecutor   else self.Executor

    @property
    def futures    (self) -> dict[object, object]  :
        return None if not self.hasFutures    else self.Futures

    @property
    def workersize (self) -> int                   :
        return None if not self.hasWorkerSize else self.WorkerSize

    @property
    def hasExecutor   (self) -> bool:
        return hasattr(self, 'Executor'  ) and isinstance(self.Executor  , ThreadPoolExecutor )

    @property
    def hasFutures    (self) -> bool:
        return hasattr(self, 'Futures'   ) and isinstance(self.Futures   , dict               )

    @property
    def hasWorkerSize (self) -> bool:
        return hasattr(self, 'WorkerSize') and isinstance(self.WorkerSize, int                )
