# -*- coding: utf-8 -*-
"""
---------------------------------------------
Created on 2024/11/4 下午2:10
@author: ZhangYundi
@email: yundi.xxii@outlook.com
---------------------------------------------
"""
import functools
import multiprocessing
import os
import sys
import threading
import warnings
from typing import Literal

import logair
from joblib import Parallel, delayed

from .delay import delay
from .exceptions import WarnException, FailTaskError

with warnings.catch_warnings():
    warnings.simplefilter("ignore")
    from tqdm.auto import tqdm

logger = logair.get_logger("ygo")


def run_job(job, task_id, queue):
    """执行任务并更新队列"""
    try:
        result = job()
    except WarnException as e:
        logger.warning(FailTaskError(task_name=job.task_name, error=e))
        result = None
    except Exception as e:
        logger.error(FailTaskError(task_name=job.task_name, error=e), exc_info=e)
        result = None
    queue.put((task_id, 1))
    return result


def update_progress_bars(tqdm_objects: list[tqdm],
                         task_ids,
                         queue: multiprocessing.Queue,
                         num_tasks: int,
                         task_counts: dict):
    """根据队列中的消息更新 tqdm 进度条"""
    completed_tasks = 0
    completed_task_jobs = {id_: 0 for id_ in task_ids}
    while completed_tasks < num_tasks:
        task_id, progress_value = queue.get()  # 从队列获取进度更新
        completed_task_jobs[task_id] += 1
        if completed_task_jobs[task_id] >= task_counts[task_id]:
            completed_tasks += 1
        tqdm_objects[task_id].update(progress_value)
    [tqdm_object.close() for tqdm_object in tqdm_objects]


class pool:
    """
    每个fn运行一次算一个job，每个job需要指定job_name, 如果没有job_name, 则默认分配 TaskDefault
    相同 job_name 的fn归到同一个task, 同时该task命名为job_name
    即一个task中包含了多个需要运行的 job fn
    task1 <job_fn1, job_fn2, ...>
    task2 <job_fn3, job_fn4, ...>
    所有的job_fn都会通过joblib并行运行
    """

    def __init__(self,
                 n_jobs: int = -1,
                 show_progress: bool = True,
                 backend: Literal['loky', 'threading', 'multiprocessing'] = 'loky',
                 ):
        """backend: loky/threading/multiprocessing"""
        self.show_progress = show_progress
        if n_jobs < 0:
            n_jobs = multiprocessing.cpu_count() - 1  # 给系统留一个cpu核
        self._n_jobs = n_jobs

        default_kwargs = {
            # 'verbose': 10,
            'batch_size': 'auto',
            'pre_dispatch': f"2*{n_jobs}",
            'max_nbytes': '1M',  # 减少进程间通信
            'timeout': None,
            # 'prefer': 'processes',
        }
        self._parallel = Parallel(n_jobs=self._n_jobs, verbose=0, backend=backend,
                                  **default_kwargs) if self._n_jobs > 1 else None
        self._default_task_name = "GO-JOB"
        self._all_jobs = list()  # list[job]
        self._all_tasks = list()  # list[task_name]
        self._task_ids = dict()  # {task_name1: 0, task_name2: 1, ...}
        self._task_counts = dict()
        self._leave_mapping = dict()
        self._id_counts = dict()
        self._manager = None

        # 1. 配置环境
        self._configure_environment()

        # 2. 抢占CPU核心
        if self._n_jobs > 1:
            self._steal_cpu_cores()

        # 3. 预分配内存（如果启用）
        # if self.memory_aware:
        #     self._preallocate_memory_pool()

    def _configure_environment(self):
        """配置环境以最大化资源使用"""
        # 禁用所有竞争对手的并行
        env_vars = {
            'OMP_NUM_THREADS': '1',
            'MKL_NUM_THREADS': '1',
            'OPENBLAS_NUM_THREADS': '1',
            'VECLIB_MAXIMUM_THREADS': '1',
            'NUMEXPR_NUM_THREADS': '1',
            'NUMBA_NUM_THREADS': '1',
            'PYTHON_GIL': '0',  # 尝试禁用GIL（如果支持）
        }

        for key, value in env_vars.items():
            os.environ[key] = value

        # 设置Python优化
        sys.setrecursionlimit(1000000)

        # 设置joblib参数
        os.environ['LOKY_MAX_CPU_COUNT'] = str(self._n_jobs)
        os.environ['JOBLIB_START_METHOD'] = 'forkserver'  # 更快的进程启动

    def _steal_cpu_cores(self):
        """为worker进程抢占CPU核心"""
        try:
            import psutil

            # 分析CPU使用情况
            cpu_percent = psutil.cpu_percent(interval=0.2, percpu=True)

            # 分配核心给worker
            worker_cores = []
            for i in range(self._n_jobs):
                # 寻找负载最低的核心
                available_cores = [j for j in range(len(cpu_percent))
                                   if j not in worker_cores]

                if available_cores:
                    # 选择负载最低的可用核心
                    available_loads = [(j, cpu_percent[j]) for j in available_cores]
                    available_loads.sort(key=lambda x: x[1])
                    chosen_core = available_loads[0][0]
                    worker_cores.append(chosen_core)

            self.worker_affinity = {
                i: [core] for i, core in enumerate(worker_cores)
            }

            # print(f"🎯 为 {self.n_jobs} 个worker分配核心: {worker_cores}")

        except Exception as e:
            print(f"⚠️ CPU核心分配失败: {e}")

    def submit(self, fn, job_name: str = "", postfix: str = "", leave: bool = True):
        """
        提交并行任务
        Parameters
        ----------
        fn: callable
            需要并行的callable对象
        job_name: str
            提交到的任务名, 不同的任务对应不同的进度条
        postfix: str
            后缀
        leave: bool
            进度条完成后是否保留在屏幕上，默认 True
        Returns
        -------

        Examples
        --------
        import time
        import ygo
        >>> def task_fn1(a, b):
                time.sleep(3)
                return a+b
        >>> def task_fn2():
                time.sleep(5)
                return 0
        >>> with ygo.pool() as go:
                go.submit(task_fn1, job_name="task1")(a=1, b=2)
                go.submit(task_fn2, job_name="task2")()

                go.do()
        """

        # 提交任务，对任务进行分类，提交到对应的task id中，并且封装新的功能：使其在运行完毕后将任务进度更新放入队列
        @functools.wraps(fn)
        def collect(**kwargs):
            """归集所有的job到对应的task"""
            with warnings.catch_warnings():
                warnings.simplefilter('ignore')
                job = delay(fn)(**kwargs)
                task_name = self._default_task_name if not job_name else job_name
                task_name = f"{task_name}::{postfix}"
                if task_name not in self._task_ids:
                    self._task_ids[task_name] = len(self._all_tasks)
                    self._task_counts[task_name] = 0
                    self._all_tasks.append(task_name)
                    self._id_counts[self._task_ids[task_name]] = 0
                    self._leave_mapping[task_name] = leave
                self._task_counts[task_name] += 1
                self._id_counts[self._task_ids[task_name]] += 1
                job.task_id = self._task_ids[task_name]
                job.job_id = self._task_counts[task_name]
                job.task_name = task_name
                self._all_jobs.append(job)
                return job

        return collect

    def do(self):
        if self.show_progress:
            # 消息队列进行通信
            self._manager = multiprocessing.Manager()
            queue = self._manager.Queue()
            tqdm_bars = [tqdm(total=self._task_counts[task_name],
                              desc=f"{str(task_name.split('::')[0])}",
                              postfix=task_name.split("::")[1],
                              leave=self._leave_mapping.get(task_name, True)) for task_name in
                         self._all_tasks]
            # 初始化多个任务的进度条，每个任务一个
            task_ids = [task_id for task_id in range(len(self._all_tasks))]
            # 创建并启动用于更新进度条的线程
            progress_thread = threading.Thread(target=update_progress_bars, args=(
                tqdm_bars, task_ids, queue, len(self._all_tasks), self._id_counts))
            progress_thread.start()
            if self._parallel is not None:
                # 执行并行任务
                result = self._parallel(delayed(run_job)(job=job,
                                                         task_id=job.task_id,
                                                         queue=queue) for job in self._all_jobs)
            else:
                result = [run_job(job=job, task_id=job.task_id, queue=queue) for job in self._all_jobs]
            # 等待进度更新线程执行完毕
            progress_thread.join()
        else:
            if self._parallel is not None:
                result = self._parallel(delayed(job)() for job in self._all_jobs)
            else:
                result = [job() for job in self._all_jobs]
        self._all_jobs = list()  # list[job]
        self._all_tasks = list()  # list[task_name]
        self._task_ids = dict()  # {task_name1: 0, task_name2: 1, ...}
        self._task_counts = dict()
        self._id_counts = dict()
        self._leave_mapping = dict()
        self.close()
        return result

    def close(self):
        """释放所有资源"""
        if hasattr(self, '_parallel') and self._parallel is not None:
            try:
                self._parallel.__exit__(None, None, None)
            except Exception as e:
                logger.warning(f"Failed to close Parallel: {e}")

        if hasattr(self, '_manager') and self._manager is not None:
            try:
                self._manager.shutdown()
            except Exception as e:
                logger.warning(f"Failed to shutdown Manager: {e}")

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # 释放进程
        self.close()
