#pragma once
#ifndef THREAD_POOL_H
#define THREAD_POOL_H

#include <vector>
#include <queue>
#include <atomic>
#include <future>
#include <stdexcept>
#include <mutex>
#include <condition_variable>
#include <thread>
#include <functional>
#include <string>

#ifdef __linux__

#elif _WIN32
#include "Windows.h"
#endif


namespace std
{
#if defined(_WIN32) || defined(_WIN64)
    // 如果 __try 和 __except 不可用，可以使用以下替代实现
    static void SetThreadNameAlternative(DWORD threadId, const char* threadName) {
        static const DWORD MS_VC_EXCEPTION = 0x406D1388;
        struct THREADNAME_INFO {
            DWORD dwType;        // 必须为 0x1000
            LPCSTR szName;       // 线程名称
            DWORD dwThreadID;    // 线程 ID
            DWORD dwFlags;       // 保留，必须为 0
        } info = { 0x1000, threadName, threadId, 0 };

        // 使用标准 C++ 异常处理
        try {
            RaiseException(MS_VC_EXCEPTION, 0, sizeof(info) / sizeof(ULONG_PTR), (ULONG_PTR*)&info);
        }
        catch (...) {
            // 忽略异常
        }
    }

    // Windows 平台设置线程名称
    static void SetThreadDescription(HANDLE hThread, const std::string& name)
    {
        PCWSTR szThreadDescription = std::wstring(name.begin(), name.end()).c_str();
        typedef LONG(NTAPI* PFN_RTL_SET_THREAD_DESCRIPTION)(HANDLE, PCWSTR);
        static PFN_RTL_SET_THREAD_DESCRIPTION pRtlSetThreadDescription = nullptr;

        if (pRtlSetThreadDescription == nullptr) {
            HMODULE hKernelBase = GetModuleHandleW(L"KernelBase.dll");
            if (hKernelBase != nullptr) {
                pRtlSetThreadDescription = (PFN_RTL_SET_THREAD_DESCRIPTION)GetProcAddress(hKernelBase, "RtlSetThreadDescription");
            }
        }

        if (pRtlSetThreadDescription != nullptr) {
            pRtlSetThreadDescription(hThread, szThreadDescription);
        }
        else
        {
            // 如果函数不可用，可以在这里添加兼容性处理代码
            // 例如，使用其他方式设置线程名称
            // 或者记录日志信息
            //OutputDebugStringW(L"SetThreadDescription function is not available.\n");
            SetThreadNameAlternative(GetCurrentThreadId(), name.c_str());
        }
    }
#endif

    // 跨平台设置线程名称
    inline void setThreadName(const std::string& name)
    {
#if defined(_WIN32) || defined(_WIN64)
        // Windows 平台设置线程名称 : win 下实测直接运行可执行程序稳定crash，通过win事件查看器捕获，暂时屏蔽win平台的线程名设置
        //SetThreadDescription(GetCurrentThread(), name);
        (void)name; // 避免未使用变量的警告
#elif defined(__linux__) || defined(__unix__)
        // Linux 平台设置线程名称
        pthread_setname_np(pthread_self(), name.c_str());
#else
        // 其他平台不支持
        (void)name; // 避免未使用变量的警告
#endif
    }

    // 线程池最大容量,应尽量设小一点
#define THREADPOOL_MAX_NUM 8
// 线程池是否可以自动增长(如果需要,且不超过 THREADPOOL_MAX_NUM)
//#define THREADPOOL_AUTO_GROW

    // 线程池,可以提交变参函数或拉姆达表达式的匿名函数执行,可以获取执行返回值
    // 不直接支持类成员函数, 支持类静态成员函数或全局函数,Opteron()函数等
    class threadpool
    {
        unsigned short _initSize;       // 初始化线程数量
        using Task = std::function<void()>; // 定义类型
        std::vector<std::thread> _pool;          // 线程池
        std::queue<Task> _tasks;            // 任务队列
        std::mutex _lock;                   // 任务队列同步锁
#ifdef THREADPOOL_AUTO_GROW
        std::mutex _lockGrow;               // 线程池增长同步锁
#endif // !THREADPOOL_AUTO_GROW
        std::condition_variable _task_cv;   // 条件阻塞
        std::atomic<bool> _run{ true };     // 线程池是否执行
        std::atomic<int>  _idlThrNum{ 0 };  // 空闲线程数量

    public:
        inline threadpool(unsigned short size = 4) { _initSize = size; addThread(size); }
        inline ~threadpool()
        {
            _run = false;
            _task_cv.notify_all(); // 唤醒所有线程执行
            for (std::thread& thread : _pool) {
                if (thread.joinable())
                    thread.join(); // 等待任务结束
            }
        }

    public:
        // 提交一个任务
        template<class F, class... Args>
        auto commit(F&& f, Args&&... args, const std::string& threadName = "") -> std::future<decltype(f(args...))>
        {
            if (!_run)    // stoped ??
                throw std::runtime_error("commit on ThreadPool is stopped.");

            using RetType = decltype(f(args...)); // typename std::result_of<F(Args...)>::type, 函数 f 的返回值类型
            auto task = std::make_shared<std::packaged_task<RetType()>>(std::bind(std::forward<F>(f), std::forward<Args>(args)...));
            std::future<RetType> future = task->get_future();
            {
                std::lock_guard<std::mutex> lock{ _lock };
                _tasks.emplace([task]() {
                    (*task)();
                    });
            }
#ifdef THREADPOOL_AUTO_GROW
            if (_idlThrNum < 1 && _pool.size() < THREADPOOL_MAX_NUM)
                addThread(1, threadName);
#else
            if (!threadName.empty() && _pool.size() < THREADPOOL_MAX_NUM) {
                addThread(1, threadName);
            }
#endif // !THREADPOOL_AUTO_GROW
            _task_cv.notify_one(); // 唤醒一个线程执行

            return future;
        }

        // 提交一个无参任务, 且无返回值
        template <class F>
        void commit2(F&& task, const std::string& threadName = "")
        {
            if (!_run) return;
            {
                std::lock_guard<std::mutex> lock{ _lock };
                _tasks.emplace(std::forward<F>(task));
            }
#ifdef THREADPOOL_AUTO_GROW
            if (_idlThrNum < 1 && _pool.size() < THREADPOOL_MAX_NUM)
                addThread(1, threadName);
#else
            if (!threadName.empty() && _pool.size() < THREADPOOL_MAX_NUM) {
                addThread(1, threadName);
            }
#endif // !THREADPOOL_AUTO_GROW
            _task_cv.notify_one();
        }

        // 空闲线程数量
        int idlCount() { return _idlThrNum; }
        // 线程数量
        size_t thrCount() { return _pool.size(); }

#ifndef THREADPOOL_AUTO_GROW
    private:
#endif // !THREADPOOL_AUTO_GROW
        // 添加指定数量的线程
        void addThread(unsigned short size, const std::string& threadName = "")
        {
#ifdef THREADPOOL_AUTO_GROW
            if (!_run)    // stoped ??
                throw std::runtime_error("Grow on ThreadPool is stopped.");
            std::unique_lock<std::mutex> lockGrow{ _lockGrow }; // 自动增长锁
#endif // !THREADPOOL_AUTO_GROW
            for (; _pool.size() < THREADPOOL_MAX_NUM && size > 0; --size)
            {   // 增加线程数量,但不超过预定义数量 THREADPOOL_MAX_NUM
                _pool.emplace_back([this, threadName] {
                    std::string finalThreadName = threadName.empty() ? "PoolThd-" + std::to_string(_pool.size()) : threadName;
                    setThreadName(finalThreadName);
                    while (true) // 防止 _run==false 时立即结束,此时任务队列可能不为空
                    {
                        Task task; // 获取一个待执行的 task
                        {
                            std::unique_lock<std::mutex> lock{ _lock };
                            _task_cv.wait(lock, [this] { // wait 直到有 task, 或需要停止
                                return !_run || !_tasks.empty();
                                });
                            if (!_run && _tasks.empty())
                                return;
                            _idlThrNum--;
                            task = std::move(_tasks.front()); // 按先进先出从队列取一个 task
                            _tasks.pop();
                        }
                        task(); // 执行任务
#ifdef THREADPOOL_AUTO_GROW
                        if (_idlThrNum > 0 && _pool.size() > _initSize) // 支持自动释放空闲线程,避免峰值过后大量空闲线程
                            return;
#endif // !THREADPOOL_AUTO_GROW
                        {
                            std::unique_lock<std::mutex> lock{ _lock };
                            _idlThrNum++;
                        }
                    }
                    });
                {
                    std::unique_lock<std::mutex> lock{ _lock };
                    _idlThrNum++;
                }
            }
        }
    };

}

#endif