#include <atomic>
#include <array>
#include <cstddef>
#include <cstdint>
#include <gtest/gtest.h>
#include <memory_resource>
#include <stdexcept>
#include <thread>
#include <vector>
#include "../Renderive/base/Memory.h"
namespace {
struct Callback_State {
    std::atomic<std::uint64_t> allocation_count{};
    std::atomic<std::uint64_t> deallocation_count{};
    std::atomic_bool alignment_error{};
    std::atomic_bool deallocated_off_main_thread{};
    std::thread::id main_thread;
};
Callback_State callback_state;
void* callback_allocate(void* context, std::size_t size, std::size_t alignment) {
    auto* state = static_cast<Callback_State*>(context);
    void* result{};
    try {
        result = std::pmr::new_delete_resource()->allocate(size, alignment);
    }
    catch (...) {
        return nullptr;
    }
    state->allocation_count.fetch_add(1, std::memory_order_relaxed);
    if (reinterpret_cast<std::uintptr_t>(result) % alignment)
        state->alignment_error.store(true, std::memory_order_relaxed);
    return result;
}
void callback_deallocate(void* context, void* address, std::size_t size, std::size_t alignment) {
    auto* state = static_cast<Callback_State*>(context);
    state->deallocation_count.fetch_add(1, std::memory_order_relaxed);
    if (std::this_thread::get_id() != state->main_thread)
        state->deallocated_off_main_thread.store(true, std::memory_order_relaxed);
    std::pmr::new_delete_resource()->deallocate(address, size, alignment);
}
struct Payload {
    std::uint64_t value{};
};
struct Large_Payload {
    std::array<std::byte, 128 * 1024> data{};
};
}
TEST(Memory_Callbacks, ValidatesCallbacksAndTracksDeallocationThread) {
    callback_state.main_thread = std::this_thread::get_id();
    renderive::Memory_Callbacks invalid_callbacks;
    invalid_callbacks.context = &callback_state;
    invalid_callbacks.allocate = callback_allocate;
    EXPECT_THROW(renderive::set_memory_callbacks(invalid_callbacks), std::invalid_argument);
    invalid_callbacks.allocate = nullptr;
    invalid_callbacks.deallocate = callback_deallocate;
    EXPECT_THROW(renderive::set_memory_callbacks(invalid_callbacks), std::invalid_argument);
    renderive::Memory_Callbacks callbacks;
    callbacks.context = &callback_state;
    callbacks.allocate = callback_allocate;
    callbacks.deallocate = callback_deallocate;
    renderive::set_memory_callbacks(callbacks);
    EXPECT_THROW(renderive::set_memory_callbacks(callbacks), std::logic_error);
    {
        std::pmr::vector<std::uint64_t> values(renderive::memory_resource());
        values.resize(8192, 11);
        auto owner = renderive::make_shared<Payload>();
        auto large_owner = renderive::make_shared<Large_Payload>();
        owner->value = values.front();
        std::thread release_thread([owner = std::move(owner), large_owner = std::move(large_owner)]() mutable {
            owner.reset();
            large_owner.reset();
        });
        release_thread.join();
    }
    renderive::release_unused_memory();
    renderive::Memory_Stats stats = renderive::memory_stats();
    EXPECT_FALSE(callback_state.alignment_error.load(std::memory_order_relaxed));
    EXPECT_GT(callback_state.allocation_count.load(std::memory_order_relaxed), 0);
    EXPECT_EQ(stats.upstream_allocation_count, callback_state.allocation_count.load(std::memory_order_relaxed));
    EXPECT_GT(callback_state.deallocation_count.load(std::memory_order_relaxed), 0);
    EXPECT_TRUE(callback_state.deallocated_off_main_thread.load(std::memory_order_relaxed));
    renderive::shutdown_memory();
    EXPECT_EQ(renderive::memory_stats().upstream_current_bytes, 0);
}
int main(int argc, char** argv) {
    testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}