#include "thread_pool.h"

#include <atomic>
#include <gtest/gtest.h>
#include <latch>
#include <memory>
#include <limits>
#include <stdexcept>
#include <string>

namespace
{

void increment(std::atomic<int>& count)
{
    ++count;
}

void block(std::shared_future<void> release, std::latch& started,
           std::atomic<int>& completed)
{
    started.count_down();
    release.wait();
    ++completed;
}

void fail()
{
    throw std::runtime_error("Task failed");
}

void assignValue(const std::unique_ptr<int>& value, std::atomic<int>& result)
{
    result = *value;
}

TEST(ThreadPool, CapacityAndReuse)
{
    std::atomic<int> completed{0};
    std::latch started(2);
    auto pool = ThreadPool::create(2).value();
    // Promise destruction releases waiters after a failed assertion.
    std::promise<void> release;
    const auto ready = release.get_future().share();
    auto first = pool->trySubmit(
        std::bind_front(block, ready, std::ref(started), std::ref(completed)))
            .value();
    auto second = pool->trySubmit(
        std::bind_front(block, ready, std::ref(started), std::ref(completed)))
            .value();
    ASSERT_TRUE(first.has_value() && second.has_value())
        << "Accept up to capacity";
    started.wait();
    auto full = pool->trySubmit(
        std::bind_front(increment, std::ref(completed)));
    ASSERT_TRUE(full.has_value() && !*full) << "Full pool is not an error";
    release.set_value();
    first->get();
    second->get();
    pool->waitIdle();
    EXPECT_EQ(completed.load(), 2) << "Rejected task never executes";
    auto next = pool->trySubmit(
        std::bind_front(increment, std::ref(completed))).value();
    ASSERT_TRUE(next.has_value()) << "Reuse a worker";
    next->get();
    pool->waitIdle();
    EXPECT_EQ(completed.load(), 3) << "Execute reused worker task";
}

TEST(ThreadPool, MoveOnlyTasks)
{
    std::atomic<int> result{0};
    auto pool = ThreadPool::create(1).value();
    auto next = pool->trySubmit(std::bind_front(
        assignValue, std::make_unique<int>(42), std::ref(result))).value();
    ASSERT_TRUE(next.has_value()) << "Accept move-only task";
    next->get();
    pool->waitIdle();
    EXPECT_EQ(result.load(), 42) << "Execute move-only task";
}

void submitBlocked(ThreadPool& pool, std::shared_future<void> release,
                   std::latch& started, std::atomic<int>& completed,
                   std::atomic<int>& accepted)
{
    auto submitted = pool.trySubmit(std::bind_front(
        block, release, std::ref(started), std::ref(completed)));
    if(submitted && *submitted)
    {
        ++accepted;
    }
}

TEST(ThreadPool, ConcurrentSubmission)
{
    std::atomic<int> completed{0};
    std::atomic<int> accepted{0};
    std::latch started(4);
    auto pool = ThreadPool::create(4).value();
    std::promise<void> release;
    const auto ready = release.get_future().share();
    std::vector<std::jthread> producers;
    for(int i = 0; i < 16; ++i)
    {
        producers.emplace_back(submitBlocked, std::ref(*pool), ready,
                               std::ref(started), std::ref(completed),
                               std::ref(accepted));
    }
    producers.clear();
    EXPECT_EQ(accepted.load(), 4) << "Concurrent submissions respect capacity";
    started.wait();
    release.set_value();
    pool->waitIdle();
    EXPECT_EQ(completed.load(), 4) << "Execute each accepted task once";
}

TEST(ThreadPool, Destruction)
{
    std::atomic<int> completed{0};
    std::latch started(2);
    std::optional<std::future<void>> first;
    std::optional<std::future<void>> second;
    {
        auto pool = ThreadPool::create(2).value();
        std::promise<void> release;
        const auto ready = release.get_future().share();
        first = pool->trySubmit(std::bind_front(
            block, ready, std::ref(started), std::ref(completed))).value();
        second = pool->trySubmit(std::bind_front(
            block, ready, std::ref(started), std::ref(completed))).value();
        ASSERT_TRUE(first.has_value() && second.has_value()) << "Accept tasks";
        // Promise destruction releases work immediately before pool teardown.
    }
    first->get();
    second->get();
    EXPECT_EQ(completed.load(), 2) << "Destruction finishes all accepted work";
}

TEST(ThreadPool, InvalidArguments)
{
    auto invalid = ThreadPool::create(0);
    ASSERT_TRUE(!invalid) << "Reject zero capacity";
    ASSERT_TRUE(!invalid.error().msg().empty()) << "Explain invalid capacity";
    EXPECT_THROW(ThreadPool::create(
        std::numeric_limits<std::size_t>::max()), std::length_error);
    auto pool = ThreadPool::create(1).value();
    auto empty = pool->trySubmit({});
    ASSERT_TRUE(!empty) << "Reject empty task with an error";
    ASSERT_TRUE(!empty.error().msg().empty()) << "Explain invalid task";
    std::atomic<int> completed{0};
    auto valid = pool->trySubmit(
        std::bind_front(increment, std::ref(completed)));
    ASSERT_TRUE(valid && *valid)
        << "Invalid submission leaves capacity available";
    (*valid)->get();
    pool->waitIdle();
}

}

TEST(ThreadPoolDeathTest, UnexpectedTaskExceptionTerminates)
{
    ASSERT_DEATH({
        auto pool = ThreadPool::create(1).value();
        auto submitted = pool->trySubmit(fail).value();
        submitted->get();
    }, "");
}