#include <gtest/gtest.h>

#include <algorithm>
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstdint>
#include <memory>
#include <mutex>
#include <thread>
#include <vector>

#include "renderive/renderable/Renderable.hpp"
#include "renderive/scene/Scene.hpp"

namespace {

using namespace std::chrono_literals;

class Gate {
public:
    void arrive_and_wait() {
        std::unique_lock lock(mutex_);
        arrived_ = true;
        condition_.notify_all();
        condition_.wait(lock, [this] { return open_; });
    }

    void wait_until_arrived() {
        std::unique_lock lock(mutex_);
        ASSERT_TRUE(condition_.wait_for(lock, 2s, [this] { return arrived_; }));
    }

    void open() {
        std::lock_guard lock(mutex_);
        open_ = true;
        condition_.notify_all();
    }

private:
    std::mutex mutex_;
    std::condition_variable condition_;
    bool arrived_{};
    bool open_{};
};

template <class Predicate>
bool wait_until(Predicate&& predicate) {
    const auto deadline = std::chrono::steady_clock::now() + 2s;
    while (!predicate()) {
        if (std::chrono::steady_clock::now() >= deadline)
            return false;
        std::this_thread::yield();
    }
    return true;
}

class Counting_Renderable : public Renderable_Base {
public:
    explicit Counting_Renderable(Scene_Base& scene)
        : Renderable_Base(scene, {.cache_enabled = false}) {}

    std::atomic<int> prepare_count{};

protected:
    void prepare(const Prepare_Render_Context&) override {
        prepare_count.fetch_add(1, std::memory_order_release);
    }
};

class Blocking_Renderable : public Counting_Renderable {
public:
    Blocking_Renderable(Scene_Base& scene, std::shared_ptr<Gate> gate,
                        int block_on_call = 1,
                        std::atomic<bool>* destroyed = nullptr)
        : Counting_Renderable(scene), gate_(std::move(gate)),
          block_on_call_(block_on_call), destroyed_(destroyed) {}

    ~Blocking_Renderable() override {
        if (destroyed_)
            destroyed_->store(true, std::memory_order_release);
    }

protected:
    void prepare(const Prepare_Render_Context&) override {
        const int call = prepare_count.fetch_add(1, std::memory_order_acq_rel) + 1;
        if (call == block_on_call_)
            gate_->arrive_and_wait();
    }

private:
    std::shared_ptr<Gate> gate_;
    int block_on_call_{};
    std::atomic<bool>* destroyed_{};
};

TEST(dynamic_renderable_lifecycle_test,
     attach_during_render_is_non_blocking_and_visible_on_the_next_frame) {
    Scene2D_Context<> scene;
    auto gate = std::make_shared<Gate>();
    auto current = std::make_shared<Blocking_Renderable>(scene, gate);
    auto attached = std::make_shared<Counting_Renderable>(scene);
    scene.attach_renderable(current);

    scene.render();
    gate->wait_until_arrived();
    std::atomic<bool> attach_finished{};
    std::thread controller([&] {
        scene.attach_renderable(attached);
        attach_finished.store(true, std::memory_order_release);
    });
    EXPECT_TRUE(wait_until([&] { return attach_finished.load(std::memory_order_acquire); }));
    EXPECT_EQ(attached->prepare_count.load(std::memory_order_acquire), 0);

    gate->open();
    controller.join();
    scene.wait_for_render();
    EXPECT_EQ(attached->prepare_count.load(std::memory_order_acquire), 0);

    scene.render();
    scene.wait_for_render();
    EXPECT_EQ(attached->prepare_count.load(std::memory_order_acquire), 1);
}

TEST(dynamic_renderable_lifecycle_test,
     detach_during_render_keeps_the_frame_owner_alive_until_execution_finishes) {
    Scene2D_Context<> scene;
    auto gate = std::make_shared<Gate>();
    std::atomic<bool> destroyed{};
    auto renderable = std::make_shared<Blocking_Renderable>(scene, gate, 1, &destroyed);
    std::weak_ptr<Blocking_Renderable> weak = renderable;
    scene.attach_renderable(renderable);

    scene.render();
    gate->wait_until_arrived();
    std::atomic<bool> detach_finished{};
    std::thread controller([&] {
        scene.detach_renderable(renderable);
        detach_finished.store(true, std::memory_order_release);
    });
    EXPECT_TRUE(wait_until([&] { return detach_finished.load(std::memory_order_acquire); }));
    controller.join();
    renderable.reset();
    EXPECT_FALSE(destroyed.load(std::memory_order_acquire));
    EXPECT_FALSE(weak.expired());

    gate->open();
    scene.wait_for_render();
    EXPECT_TRUE(wait_until([&] { return destroyed.load(std::memory_order_acquire); }));
    EXPECT_TRUE(weak.expired());
    EXPECT_EQ(scene.renderable_count(), 0u);
}

TEST(dynamic_renderable_lifecycle_test,
     reparent_during_render_changes_only_the_next_frame_dependency_plan) {
    Scene2D_Context<> scene;
    auto first_gate = std::make_shared<Gate>();
    auto second_gate = std::make_shared<Gate>();
    auto first_parent = std::make_shared<Blocking_Renderable>(scene, first_gate, 1);
    auto second_parent = std::make_shared<Blocking_Renderable>(scene, second_gate, 2);
    auto child = std::make_shared<Counting_Renderable>(scene);
    scene.attach_renderable(first_parent);
    scene.attach_renderable(second_parent);
    scene.attach_renderable(child, {
        .dependency_parents = {first_parent}
    });

    scene.render();
    first_gate->wait_until_arrived();
    scene.set_dependency_parent(child, second_parent);
    EXPECT_EQ(child->prepare_count.load(std::memory_order_acquire), 0);
    first_gate->open();
    scene.wait_for_render();
    EXPECT_EQ(child->prepare_count.load(std::memory_order_acquire), 1);

    scene.render();
    second_gate->wait_until_arrived();
    EXPECT_EQ(child->prepare_count.load(std::memory_order_acquire), 1);
    second_gate->open();
    scene.wait_for_render();
    EXPECT_EQ(child->prepare_count.load(std::memory_order_acquire), 2);
}

class Rebuildable_Renderable final : public Renderable_Base {
public:
    Rebuildable_Renderable(Scene_Base& scene, std::shared_ptr<Gate> first_graph_gate)
        : Renderable_Base(scene, {.cache_enabled = false}),
          first_graph_gate_(std::move(first_graph_gate)) {}

    void publish_second_graph() {
        graph_version_.store(2, std::memory_order_release);
        rebuild_render_graph();
    }

    std::vector<int> executions() const {
        std::lock_guard lock(executions_mutex_);
        return executions_;
    }

protected:
    void build_prepare_graph(Renderable_Graph_Builder& builder) override {
        const int version = graph_version_.load(std::memory_order_acquire);
        builder.emplace("versioned_prepare", version == 1 ? "Graph A" : "Graph B",
                        [this, version](const Prepare_Render_Context&) {
            if (version == 1)
                first_graph_gate_->arrive_and_wait();
            std::lock_guard lock(executions_mutex_);
            executions_.push_back(version);
        });
    }

private:
    std::shared_ptr<Gate> first_graph_gate_;
    std::atomic<int> graph_version_{1};
    mutable std::mutex executions_mutex_;
    std::vector<int> executions_;
};

TEST(dynamic_renderable_lifecycle_test,
     graph_rebuild_during_render_publishes_an_immutable_graph_for_the_next_frame) {
    Scene2D_Context<> scene;
    auto gate = std::make_shared<Gate>();
    auto renderable = std::make_shared<Rebuildable_Renderable>(scene, gate);
    scene.attach_renderable(renderable);

    scene.render();
    gate->wait_until_arrived();
    renderable->publish_second_graph();
    gate->open();
    scene.wait_for_render();
    EXPECT_EQ(renderable->executions(), std::vector<int>({1}));

    scene.render();
    scene.wait_for_render();
    EXPECT_EQ(renderable->executions(), std::vector<int>({1, 2}));
}

class Self_Detaching_Renderable final : public Counting_Renderable {
public:
    using Counting_Renderable::Counting_Renderable;

    std::weak_ptr<Self_Detaching_Renderable> self;

protected:
    void prepare(const Prepare_Render_Context&) override {
        prepare_count.fetch_add(1, std::memory_order_release);
        if (auto owner = self.lock())
            scene().detach_renderable(owner);
    }
};

TEST(dynamic_renderable_lifecycle_test,
     renderable_can_detach_itself_without_deadlock_and_finishes_the_current_frame) {
    Scene2D_Context<> scene;
    auto renderable = std::make_shared<Self_Detaching_Renderable>(scene);
    renderable->self = renderable;
    scene.attach_renderable(renderable);

    scene.render();
    scene.wait_for_render();
    EXPECT_EQ(renderable->prepare_count.load(std::memory_order_acquire), 1);
    EXPECT_EQ(scene.renderable_count(), 0u);

    scene.render();
    scene.wait_for_render();
    EXPECT_EQ(renderable->prepare_count.load(std::memory_order_acquire), 1);
}

class Tracking_Color_Cache final : public Color_Cache {
public:
    Tracking_Color_Cache() : instance_id(next_instance_id.fetch_add(1) + 1) {}
    explicit Tracking_Color_Cache(std::pmr::memory_resource&) : Tracking_Color_Cache() {}
    ~Tracking_Color_Cache() override {
        if (instance_id == watched_instance.load(std::memory_order_acquire))
            watched_destroyed.store(true, std::memory_order_release);
    }

    void clear() override {
        samples.clear();
    }

    void composite(const Color_Cache& source) override {
        const auto& cache = dynamic_cast<const Tracking_Color_Cache&>(source);
        samples.insert(samples.end(), cache.samples.begin(), cache.samples.end());
    }

    static void reset() {
        next_instance_id.store(0, std::memory_order_release);
        watched_instance.store(0, std::memory_order_release);
        watched_destroyed.store(false, std::memory_order_release);
    }

    int instance_id{};
    std::vector<int> samples;
    inline static std::atomic<int> next_instance_id{};
    inline static std::atomic<int> watched_instance{};
    inline static std::atomic<bool> watched_destroyed{};
};

class Blocking_Painted_Renderable final : public Renderable_Base {
public:
    Blocking_Painted_Renderable(Scene_Base& scene, std::shared_ptr<Gate> gate)
        : Renderable_Base(scene, {.cache_enabled = false}), gate_(std::move(gate)) {}

protected:
    void prepare(const Prepare_Render_Context&) override {
        if (prepare_count_.fetch_add(1, std::memory_order_acq_rel) == 0)
            gate_->arrive_and_wait();
    }

    void build_paint_graph(Renderable_Graph_Builder& builder) override {
        builder.emplace("paint", "Paint",
                        [](const Paint_Render_Context& context) {
            auto& cache = dynamic_cast<Tracking_Color_Cache&>(context.color_cache);
            cache.samples.push_back(cache.instance_id);
        });
    }

private:
    std::shared_ptr<Gate> gate_;
    std::atomic<int> prepare_count_{};
};

TEST(dynamic_renderable_lifecycle_test,
     detached_color_cache_lives_through_the_frame_and_reattach_uses_a_new_cache) {
    Tracking_Color_Cache::reset();
    Scene2D_Context<Low_Latency_Strategy<Scene2D_Frame_Data>, Tracking_Color_Cache> scene;
    auto gate = std::make_shared<Gate>();
    auto renderable = std::make_shared<Blocking_Painted_Renderable>(scene, gate);
    scene.attach_renderable(renderable);
    ASSERT_EQ(Tracking_Color_Cache::next_instance_id.load(), 2);
    Tracking_Color_Cache::watched_instance.store(2, std::memory_order_release);

    scene.render();
    gate->wait_until_arrived();
    scene.detach_renderable(renderable);
    EXPECT_FALSE(Tracking_Color_Cache::watched_destroyed.load(std::memory_order_acquire));
    scene.attach_renderable(renderable);
    ASSERT_EQ(Tracking_Color_Cache::next_instance_id.load(), 3);
    gate->open();
    scene.wait_for_render();
    EXPECT_TRUE(wait_until([] {
        return Tracking_Color_Cache::watched_destroyed.load(std::memory_order_acquire);
    }));
    scene.with_final_color_cache([](const Tracking_Color_Cache& cache) {
        EXPECT_EQ(cache.samples, std::vector<int>({2}));
    });

    scene.render();
    scene.wait_for_render();
    scene.with_final_color_cache([](const Tracking_Color_Cache& cache) {
        EXPECT_EQ(cache.samples, std::vector<int>({3}));
    });
}

class Sequence_Parent final : public Renderable_Base {
public:
    Sequence_Parent(Scene_Base& scene, std::atomic<std::uint64_t>& completed_sequence)
        : Renderable_Base(scene, {.cache_enabled = false}),
          completed_sequence_(completed_sequence) {}

protected:
    void prepare(const Prepare_Render_Context& context) override {
        completed_sequence_.store(context.frame.render_sequence, std::memory_order_release);
    }

private:
    std::atomic<std::uint64_t>& completed_sequence_;
};

class Sequence_Child final : public Renderable_Base {
public:
    Sequence_Child(Scene_Base& scene, std::atomic<std::uint64_t>& parent_sequence,
                   std::atomic<int>& violations)
        : Renderable_Base(scene, {.cache_enabled = false}),
          parent_sequence_(parent_sequence), violations_(violations) {}

    void invalidate_graph() {
        rebuild_render_graph();
    }

protected:
    void prepare(const Prepare_Render_Context& context) override {
        if (parent_sequence_.load(std::memory_order_acquire) != context.frame.render_sequence)
            violations_.fetch_add(1, std::memory_order_relaxed);
    }

private:
    std::atomic<std::uint64_t>& parent_sequence_;
    std::atomic<int>& violations_;
};

TEST(dynamic_renderable_lifecycle_test,
     atomic_attach_and_concurrent_membership_updates_never_publish_partial_topology) {
    Scene2D_Context<> scene;
    std::atomic<std::uint64_t> parent_sequence{};
    std::atomic<int> violations{};
    std::atomic<int> failures{};
    auto parent = std::make_shared<Sequence_Parent>(scene, parent_sequence);
    auto alternate_parent =
        std::make_shared<Sequence_Parent>(scene, parent_sequence);
    scene.attach_renderable(parent);
    scene.attach_renderable(alternate_parent, {
        .display_parents = {parent},
        .dependency_parents = {parent}
    });

    std::atomic<bool> controls_done{};
    std::thread renderer([&] {
        while (!controls_done.load(std::memory_order_acquire)) {
            try {
                scene.render();
                scene.wait_for_render();
            } catch (...) {
                failures.fetch_add(1, std::memory_order_relaxed);
            }
        }
    });

    constexpr int controller_count = 2;
    constexpr int iterations = 1000;
    std::vector<std::thread> controllers;
    for (int controller = 0; controller < controller_count; ++controller) {
        controllers.emplace_back([&] {
            for (int index = 0; index < iterations; ++index) {
                try {
                    auto child = std::make_shared<Sequence_Child>(
                        scene, parent_sequence, violations);
                    scene.attach_renderable(child, {
                        .display_parents = {parent},
                        .dependency_parents = {parent}
                    });
                    scene.set_display_parent(child, alternate_parent);
                    scene.set_dependency_parent(child, alternate_parent);
                    child->invalidate_graph();
                    scene.set_display_parent(child, parent);
                    scene.set_dependency_parent(child, parent);
                    scene.detach_renderable(child);
                } catch (...) {
                    failures.fetch_add(1, std::memory_order_relaxed);
                }
            }
        });
    }
    for (auto& controller : controllers)
        controller.join();
    controls_done.store(true, std::memory_order_release);
    renderer.join();

    EXPECT_EQ(failures.load(std::memory_order_acquire), 0);
    EXPECT_EQ(violations.load(std::memory_order_acquire), 0);
    EXPECT_EQ(scene.renderable_count(), 2u);
}

TEST(dynamic_renderable_lifecycle_test,
     captures_before_and_after_dynamic_membership_keep_both_plan_versions_resolvable) {
    Scene2D_Context<> scene;
    auto first = std::make_shared<Counting_Renderable>(scene);
    auto second = std::make_shared<Counting_Renderable>(scene);
    scene.attach_renderable(first);
    const Capture_Session_Id session_id = scene.capture_frames(2);

    scene.render();
    scene.wait_for_render();
    scene.attach_renderable(second);
    scene.render();
    scene.wait_for_render();

    const auto session = scene.capture_session(session_id);
    ASSERT_TRUE(session);
    ASSERT_EQ(session->frames.size(), 2u);
    const auto first_version = session->frames[0].snapshot->render_plan_version;
    const auto second_version = session->frames[1].snapshot->render_plan_version;
    EXPECT_NE(first_version, second_version);
    const auto first_plan = scene.find_render_plan(first_version);
    const auto second_plan = scene.find_render_plan(second_version);
    ASSERT_TRUE(first_plan);
    ASSERT_TRUE(second_plan);
    EXPECT_TRUE(std::ranges::none_of(first_plan->graph.nodes, [&](const Render_Node& node) {
        return node.owner_id == second->renderable_id();
    }));
    EXPECT_TRUE(std::ranges::any_of(second_plan->graph.nodes, [&](const Render_Node& node) {
        return node.owner_id == second->renderable_id();
    }));
}

} // namespace