#include <yandex/maps/wiki/threadutils/threadpool.h>

#include <maps/libs/common/include/exception.h>

#include <queue>
#include <mutex>
#include <thread>
#include <condition_variable>

namespace maps::wiki {

class ThreadPool::Impl
{
public:
    using Functor = std::function<void ()>;

    explicit Impl(size_t threads)
    {
        ASSERT(threads);
        while (threads--) {
            threads_.emplace_back(std::bind(&Impl::thread, this));
        }
    }

    ~Impl()
    {
        stop();
        for (auto& thread : threads_) {
            thread.join();
        }
    }

    bool push(Functor f)
    {
        std::lock_guard<std::mutex> lock(mutex_);
        if (stopped_ || shutdown_) {
            return false;
        }
        pendingTasks_.push(std::move(f));
        cond_.notify_one();
        return true;
    }

    size_t threadsCount() const
    {
        return threads_.size();
    }

    size_t activeTasksCount() const
    {
        std::lock_guard<std::mutex> lock(mutex_);
        return activeTasksCount_;
    }

    size_t pendingTasksCount() const
    {
        std::lock_guard<std::mutex> lock(mutex_);
        return pendingTasks_.size();
    }

    void shutdown()
    {
        std::unique_lock<std::mutex> lock(mutex_);
        shutdown_ = true;
        while (!pendingTasks_.empty() && !stopped_) {
            shutdownCond_.wait(lock);
        }
        stopped_ = true;
        cond_.notify_all();

        while (activeTasksCount_ > 0) {
            shutdownCond_.wait(lock);
        }
    }

    void stop()
    {
        std::lock_guard<std::mutex> lock(mutex_);
        stopped_ = true;
        cond_.notify_all();
        shutdownCond_.notify_all();
    }

private:
    class ScopedUnlocker
    {
    public:
        ScopedUnlocker(
                std::unique_lock<std::mutex>& lock,
                size_t& activeCount,
                std::condition_variable& cond)
            : lock_(lock)
            , activeCount_(activeCount)
            , shutdownCond_(cond)
        {
            ++activeCount_;
            lock_.unlock();
        }

        ~ScopedUnlocker()
        {
            lock_.lock();
            --activeCount_;
            if (activeCount_ == 0) {
                shutdownCond_.notify_all();
            }
        }
    private:
        std::unique_lock<std::mutex>& lock_;
        size_t& activeCount_;
        std::condition_variable& shutdownCond_;
    };

    void thread()
    {
        std::unique_lock<std::mutex> lock(mutex_);
        while (!stopped_) {
            try {
                while (pendingTasks_.empty()) {
                    if (shutdown_) {
                        shutdownCond_.notify_all();
                        return;
                    }
                    cond_.wait(lock);
                    if (stopped_) {
                        return;
                    }
                }
                Functor f = std::move(pendingTasks_.front());
                pendingTasks_.pop();

                ScopedUnlocker unlocker(lock, activeTasksCount_, shutdownCond_);
                f();
            }
            catch (...) {
            }
        }
    }

    size_t activeTasksCount_{0};
    bool stopped_{false};
    bool shutdown_{false};
    std::queue<Functor> pendingTasks_;
    mutable std::mutex mutex_;
    std::condition_variable cond_;
    std::condition_variable shutdownCond_;
    std::vector<std::thread> threads_;
};



ThreadPool::ThreadPool(size_t threads)
    : impl_(new Impl(threads))
{
}

ThreadPool::ThreadPool(ThreadPool&& other) noexcept
    : impl_(std::move(other.impl_))
{
}

ThreadPool::~ThreadPool() = default;

bool
ThreadPool::push(Functor f)
{
    return impl_->push(std::move(f));
}

size_t
ThreadPool::threadsCount() const
{
    return impl_->threadsCount();
}

size_t
ThreadPool::activeTasksCount() const
{
    return impl_->activeTasksCount();
}

size_t
ThreadPool::pendingTasksCount() const
{
    return impl_->pendingTasksCount();
}

void
ThreadPool::shutdown()
{
    impl_->shutdown();
}

void
ThreadPool::stop()
{
    impl_->stop();
}

} // namespace maps::wiki
