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

// ThreadPool.hpp
// Lightweight thread pool -- threads are launched at initialization and reused
// across operations to avoid thread creation/destruction cost.

#pragma once

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

namespace sangi {

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;

    // Submit a task and return a 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)(); });

            // If there are no idle threads, grow the pool (upper bound: hardware_concurrency)
            unsigned total = static_cast<unsigned>(workers_.size());
            unsigned hw = std::thread::hardware_concurrency();
            unsigned maxThreads = (hw > 1) ? hw : 4;
            if (busy_.load(std::memory_order_relaxed) >= total && total < maxThreads) {
                unsigned toAdd = std::min(total, maxThreads - total);
                for (unsigned i = 0; i < toAdd; ++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_;
};

// Global thread pool (one per process)
inline ThreadPool& threadPool() {
    static ThreadPool pool;
    return pool;
}

} // namespace sangi