#ifndef PKG_TASK_SCHEDULER_H
#define PKG_TASK_SCHEDULER_H

#include <cstdint>
#include <functional>
#include <future>
#include <mutex>
#include <condition_variable>
#include <queue>
#include <thread>
#include <vector>

#include "Core/Result.h"

namespace PKG
{
    // ─── 计数信号量 ─────────────────────────────────────────────────
    // 用于限制并发重操作（如 .tex 解码）的数量，降低峰值内存
    class Semaphore
    {
    public:
        explicit Semaphore(int count) : m_Count(count) {}

        // 获取一个许可（计数为 0 时阻塞等待）
        void Acquire()
        {
            std::unique_lock<std::mutex> lock(m_Mutex);
            m_Cond.wait(lock, [this] { return m_Count > 0; });
            m_Count--;
        }

        // 释放一个许可（唤醒一个等待者）
        void Release()
        {
            std::lock_guard<std::mutex> lock(m_Mutex);
            m_Count++;
            m_Cond.notify_one();
        }

    private:
        std::mutex            m_Mutex;
        std::condition_variable m_Cond;
        int                   m_Count;
    };

    // ─── RAII 信号量守卫 ────────────────────────────────────────────
    // 构造时获取信号量，析构时自动释放（异常安全）
    class SemaphoreGuard
    {
    public:
        explicit SemaphoreGuard(Semaphore& sem) : m_Sem(sem) { m_Sem.Acquire(); }
        ~SemaphoreGuard() { m_Sem.Release(); }

        SemaphoreGuard(const SemaphoreGuard&) = delete;
        SemaphoreGuard& operator=(const SemaphoreGuard&) = delete;

    private:
        Semaphore& m_Sem;
    };

    // ─── 线程池 ─────────────────────────────────────────────────────
    // 固定大小的线程池，提交任务返回 future 用于获取结果
    // 支持 WaitAll() 等待所有已提交任务完成
    class TaskScheduler
    {
    public:
        // threadCount = 0 时自动使用 hardware_concurrency
        explicit TaskScheduler(uint32_t threadCount = 0);
        ~TaskScheduler();

        TaskScheduler(const TaskScheduler&) = delete;
        TaskScheduler& operator=(const TaskScheduler&) = delete;

        // 提交任务，返回 future 以获取结果或等待完成
        std::future<Result<void>> Submit(std::function<Result<void>()> task);

        // 阻塞等待所有已提交任务完成
        void WaitAll();

        uint32_t ThreadCount() const { return static_cast<uint32_t>(m_Workers.size()); }

    private:
        std::vector<std::thread>          m_Workers;        // 工作线程
        std::queue<std::function<void()>> m_Tasks;          // 待执行任务队列
        std::mutex                        m_QueueMutex;     // 保护 m_Tasks
        std::condition_variable           m_Condition;      // 唤醒等待任务的工作线程
        bool                              m_Stop = false;   // 析构标志
        uint32_t                          m_ActiveTasks = 0; // 正在执行的任务数
        std::mutex                        m_ActiveMutex;    // 保护 m_ActiveTasks
        std::condition_variable           m_ActiveCondition; // WaitAll 等待条件
    };
}

#endif // PKG_TASK_SCHEDULER_H