#include <gtest/gtest.h>

#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstdint>
#include <deque>
#include <exception>
#include <functional>
#include <memory>
#include <mutex>
#include <stdexcept>
#include <thread>
#include <vector>

#include "renderive/render_graph/External_Operation.hpp"
#include "renderive/render_graph/Render_Plan.hpp"
#include "renderive/render_graph/detail/Render_Graph_Runtime.hpp"

namespace {

class Single_Worker_Scheduler final {
public:
    Single_Worker_Scheduler()
        : worker_([this] { run(); }) {}

    ~Single_Worker_Scheduler() {
        {
            std::lock_guard lock(mutex_);
            stopping_ = true;
        }
        ready_.notify_all();
        worker_.join();
    }

    void schedule(std::function<void()> function) {
        {
            std::lock_guard lock(mutex_);
            if (stopping_)
                throw std::logic_error("scheduler is stopping");
            queue_.push_back(std::move(function));
        }
        ready_.notify_one();
    }

private:
    void run() {
        for (;;) {
            std::function<void()> function;
            {
                std::unique_lock lock(mutex_);
                ready_.wait(lock, [this] {
                    return stopping_ || !queue_.empty();
                });
                if (stopping_ && queue_.empty())
                    return;
                function = std::move(queue_.front());
                queue_.pop_front();
            }
            function();
        }
    }

    std::mutex mutex_;
    std::condition_variable ready_;
    std::deque<std::function<void()>> queue_;
    std::thread worker_;
    bool stopping_{};
};

std::shared_ptr<const Render_Plan> two_node_plan() {
    Render_Graph_Builder builder;
    const auto submit = builder.emplace(
        301, 30, "Submit GPU", Render_Node_Kind::render);
    const auto publish = builder.emplace(
        302, 30, "Publish Frame", Render_Node_Kind::composite);
    builder.precede(submit, publish);
    Render_Plan_History history;
    return history.publish(std::move(builder).finish());
}

std::vector<Node_Execution*> execution_slots(
    const Render_Plan& plan, std::vector<Node_Execution>& storage) {
    storage.resize(plan.graph.nodes.size());
    std::vector<Node_Execution*> result(storage.size());
    for (const auto& node : plan.graph.nodes) {
        storage[node.execution_index].node_id = node.node_id;
        result[node.execution_index] = &storage[node.execution_index];
    }
    return result;
}

} // namespace

TEST(external_operation_test, completion_before_subscription_is_delivered_once) {
    External_Operation_Source source;
    const auto operation = source.operation();
    EXPECT_TRUE(source.complete());
    EXPECT_FALSE(source.complete());

    int completion_count{};
    operation.on_complete([&](std::exception_ptr error) {
        EXPECT_FALSE(error);
        ++completion_count;
    });
    EXPECT_EQ(completion_count, 1);
}

TEST(external_operation_test,
     callback_exceptions_are_contained_for_both_completion_orders) {
    External_Operation_Source completed_first;
    const auto completed_operation = completed_first.operation();
    ASSERT_TRUE(completed_first.complete());
    EXPECT_NO_THROW(completed_operation.on_complete(
        [](std::exception_ptr) { throw std::runtime_error("late callback"); }));

    External_Operation_Source subscribed_first;
    const auto subscribed_operation = subscribed_first.operation();
    subscribed_operation.on_complete(
        [](std::exception_ptr) { throw std::runtime_error("early callback"); });
    EXPECT_TRUE(subscribed_first.complete());
}

TEST(render_graph_runtime_test,
     external_successor_stays_blocked_until_operation_completes) {
    const auto plan = two_node_plan();
    Single_Worker_Scheduler scheduler;
    External_Operation_Source source;
    std::vector<Node_Execution> execution_storage;
    auto slots = execution_slots(*plan, execution_storage);
    std::atomic<bool> submit_started{};
    std::atomic<bool> publish_executed{};

    renderive::render_graph::detail::Render_Graph_Runtime runtime(
        *plan, slots,
        [&](std::size_t index, Node_Execution_Metrics*) {
            if (index == 0) {
                submit_started.store(true, std::memory_order_release);
                submit_started.notify_all();
                return Node_Execution_Result::external(source.operation());
            }
            publish_executed.store(true, std::memory_order_release);
            return Node_Execution_Result::completed();
        },
        [&](std::function<void()> function) {
            scheduler.schedule(std::move(function));
        },
        [] { return 0U; });

    std::thread execution([&] { runtime.execute(); });
    submit_started.wait(false, std::memory_order_acquire);
    EXPECT_FALSE(publish_executed.load(std::memory_order_acquire));
    EXPECT_EQ(execution_storage[0].status,
              Node_Execution_Status::waiting_external);
    EXPECT_TRUE(source.complete());
    execution.join();

    EXPECT_TRUE(publish_executed.load(std::memory_order_acquire));
    EXPECT_EQ(execution_storage[0].status, Node_Execution_Status::complete);
    EXPECT_EQ(execution_storage[1].status, Node_Execution_Status::complete);
    EXPECT_GT(execution_storage[0].cpu_end_time_ns, 0U);
    EXPECT_GE(execution_storage[0].external_end_time_ns,
              execution_storage[0].external_start_time_ns);
}

TEST(render_graph_runtime_test,
     failed_external_node_prevents_its_successor_from_running) {
    const auto plan = two_node_plan();
    Single_Worker_Scheduler scheduler;
    External_Operation_Source source;
    std::atomic<bool> submit_started{};
    std::atomic<bool> publish_executed{};
    std::exception_ptr graph_error;

    renderive::render_graph::detail::Render_Graph_Runtime runtime(
        *plan, {},
        [&](std::size_t index, Node_Execution_Metrics*) {
            if (index == 0) {
                submit_started.store(true, std::memory_order_release);
                submit_started.notify_all();
                return Node_Execution_Result::external(source.operation());
            }
            publish_executed.store(true, std::memory_order_release);
            return Node_Execution_Result::completed();
        },
        [&](std::function<void()> function) {
            scheduler.schedule(std::move(function));
        },
        [] { return 0U; });

    std::thread execution([&] {
        try {
            runtime.execute();
        } catch (...) {
            graph_error = std::current_exception();
        }
    });
    submit_started.wait(false, std::memory_order_acquire);
    EXPECT_TRUE(source.fail(
        std::make_exception_ptr(std::runtime_error("GPU submission failed"))));
    execution.join();

    ASSERT_TRUE(graph_error);
    EXPECT_THROW(std::rethrow_exception(graph_error), std::runtime_error);
    EXPECT_FALSE(publish_executed.load(std::memory_order_acquire));
}

TEST(render_graph_runtime_test,
     pending_external_operation_does_not_occupy_scheduler_worker) {
    Render_Graph_Builder builder;
    builder.emplace(401, 40, "GPU Fence", Render_Node_Kind::render);
    Render_Plan_History history;
    const auto plan = history.publish(std::move(builder).finish());
    Single_Worker_Scheduler scheduler;
    External_Operation_Source source;
    std::atomic<bool> external_started{};
    std::mutex probe_mutex;
    std::condition_variable probe_completed;
    bool probe_ran{};

    renderive::render_graph::detail::Render_Graph_Runtime runtime(
        *plan, {},
        [&](std::size_t, Node_Execution_Metrics*) {
            external_started.store(true, std::memory_order_release);
            external_started.notify_all();
            return Node_Execution_Result::external(source.operation());
        },
        [&](std::function<void()> function) {
            scheduler.schedule(std::move(function));
        },
        [] { return 0U; });

    std::thread execution([&] { runtime.execute(); });
    external_started.wait(false, std::memory_order_acquire);
    scheduler.schedule([&] {
        {
            std::lock_guard lock(probe_mutex);
            probe_ran = true;
        }
        probe_completed.notify_one();
    });
    {
        std::unique_lock lock(probe_mutex);
        EXPECT_TRUE(probe_completed.wait_for(
            lock, std::chrono::seconds(1), [&] { return probe_ran; }));
    }
    EXPECT_TRUE(source.complete());
    execution.join();
}