// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

// ThreadPool.hpp
// 軽量スレッドプール — 初期化時にスレッドを起動し、
// 各演算で再利用することでスレッド生成/破棄コストを回避

#pragma once

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

namespace calx {

class ThreadPool {
public:
    explicit ThreadPool(unsigned initial_threads = 0)
        : stop_(false), busy_(0)
    {
        if (initial_threads == 0) {
            unsigned hw = std::thread::hardware_concurrency();
            initial_threads = (hw > 1) ? hw - 1 : 1;
        }
        for (unsigned i = 0; i < initial_threads; ++i)
            addWorker();
    }

    ~ThreadPool() {
        {
            std::lock_guard<std::mutex> lock(mutex_);
            stop_ = true;
        }
        cv_.notify_all();
        for (auto& t : workers_)
            t.join();
    }

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

    // タスクを投入し std::future を返す
    template<typename F>
    auto submit(F&& f) -> std::future<std::invoke_result_t<std::decay_t<F>>> {
        using R = std::invoke_result_t<std::decay_t<F>>;
        auto task = std::make_shared<std::packaged_task<R()>>(std::forward<F>(f));
        std::future<R> result = task->get_future();

        {
            std::lock_guard<std::mutex> lock(mutex_);
            tasks_.emplace([task]() { (*task)(); });

            // 空きスレッドがなければプールを倍に拡張
            unsigned total = static_cast<unsigned>(workers_.size());
            if (busy_.load(std::memory_order_relaxed) >= total) {
                for (unsigned i = 0; i < total; ++i)
                    addWorker();
            }
        }
        cv_.notify_one();
        return result;
    }

private:
    void addWorker() {
        workers_.emplace_back([this]() {
            for (;;) {
                std::function<void()> task;
                {
                    std::unique_lock<std::mutex> lock(mutex_);
                    cv_.wait(lock, [this]{ return stop_ || !tasks_.empty(); });
                    if (stop_ && tasks_.empty()) return;
                    task = std::move(tasks_.front());
                    tasks_.pop();
                    busy_.fetch_add(1, std::memory_order_relaxed);
                }
                task();
                busy_.fetch_sub(1, std::memory_order_relaxed);
            }
        });
    }

    std::vector<std::thread> workers_;
    std::queue<std::function<void()>> tasks_;
    std::mutex mutex_;
    std::condition_variable cv_;
    bool stop_;
    std::atomic<unsigned> busy_;
};

// グローバルスレッドプール (プロセスあたり 1 つ)
inline ThreadPool& threadPool() {
    static ThreadPool pool;
    return pool;
}

} // namespace calx