#include <atomic>
#include <gtest/gtest.h>
#include <memory_resource>
#include <new>
#include <vector>
#include "../Renderive/base/Memory.h"
#include "../Renderive/base/Frame_Memory.h"
#include "Core/Base/RingBuffer.hpp"
namespace {
struct Callback_State {
    std::atomic_bool fail{true};
};
Callback_State callback_state;
void* callback_allocate(void* context, std::size_t size, std::size_t alignment) {
    auto* state = static_cast<Callback_State*>(context);
    if (state->fail.load(std::memory_order_relaxed))
        return nullptr;
    return std::pmr::new_delete_resource()->allocate(size, alignment);
}
void callback_deallocate(void*, void* address, std::size_t size, std::size_t alignment) {
    std::pmr::new_delete_resource()->deallocate(address, size, alignment);
}
}
TEST(Memory_Failure, ConvertsCallbackFailureToBadAlloc) {
    renderive::Memory_Callbacks callbacks;
    callbacks.context = &callback_state;
    callbacks.allocate = callback_allocate;
    callbacks.deallocate = callback_deallocate;
    callback_state.fail.store(false, std::memory_order_relaxed);
    renderive::set_memory_callbacks(callbacks);
    std::pmr::memory_resource* resource = renderive::memory_resource();
    callback_state.fail.store(true, std::memory_order_relaxed);
    bool direct_failure_caught = false;
    try {
        void* unexpected = resource->allocate(128 * 1024);
        resource->deallocate(unexpected, 128 * 1024);
    }
    catch (const std::bad_alloc&) {
        direct_failure_caught = true;
    }
    callback_state.fail.store(false, std::memory_order_relaxed);
    std::pmr::vector<int> values(resource);
    callback_state.fail.store(true, std::memory_order_relaxed);
    bool container_failure_caught = false;
    try {
        values.resize(1024 * 1024);
    }
    catch (const std::bad_alloc&) {
        container_failure_caught = true;
    }
    callback_state.fail.store(false, std::memory_order_relaxed);
    {
        renderive::Frame_Arena arena(64 * 1024);
        std::pmr::vector<double> scratch(arena.resource());
        scratch.resize(16 * 1024);
    }
    callback_state.fail.store(true, std::memory_order_relaxed);
    bool frame_failure_caught = false;
    try {
        renderive::Frame_Arena arena(1024 * 1024);
        std::pmr::vector<double> scratch(arena.resource());
        scratch.resize(1024 * 1024);
    }
    catch (const std::bad_alloc&) {
        frame_failure_caught = true;
    }
    bool ring_buffer_failure_caught = false;
    try {
        Psc::RingBuffer_ST buffer(renderive::memory_resource());
        buffer.init(1024 * 1024);
    }
    catch (const std::bad_alloc&) {
        ring_buffer_failure_caught = true;
    }
    EXPECT_TRUE(direct_failure_caught);
    EXPECT_TRUE(container_failure_caught);
    EXPECT_TRUE(frame_failure_caught);
    EXPECT_TRUE(ring_buffer_failure_caught);
}
int main(int argc, char** argv) {
    testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}