#include "Memory_Resource.h"
#include <atomic>
#include <mimalloc.h>
#include <new>
namespace {
struct Allocation_Counters {
    std::atomic<std::uint64_t> live_bytes{};
    std::atomic<std::uint64_t> peak_live_bytes{};
    std::atomic<std::uint64_t> allocated_bytes_total{};
    std::atomic<std::uint64_t> allocation_count{};
    std::atomic<std::uint64_t> deallocation_count{};
    std::atomic<std::uint64_t> pmr_live_bytes{};
    std::atomic<std::uint64_t> pmr_peak_live_bytes{};
    std::atomic<std::uint64_t> pmr_allocated_bytes_total{};
    std::atomic<std::uint64_t> pmr_allocation_count{};
    std::atomic<std::uint64_t> pmr_deallocation_count{};
};
constinit Allocation_Counters allocation_counters;
void update_peak(std::atomic<std::uint64_t>& peak, std::uint64_t value) {
    auto current = peak.load(std::memory_order_relaxed);
    while (current < value && !peak.compare_exchange_weak(current, value, std::memory_order_relaxed)) {
    }
}
void record_allocation(void* pointer, bool pmr) {
    if (pointer == nullptr)
        return;
    const auto bytes = static_cast<std::uint64_t>(mi_usable_size(pointer));
    const auto live = allocation_counters.live_bytes.fetch_add(bytes, std::memory_order_relaxed) + bytes;
    allocation_counters.allocated_bytes_total.fetch_add(bytes, std::memory_order_relaxed);
    allocation_counters.allocation_count.fetch_add(1, std::memory_order_relaxed);
    update_peak(allocation_counters.peak_live_bytes, live);
    if (!pmr)
        return;
    const auto pmr_live = allocation_counters.pmr_live_bytes.fetch_add(bytes, std::memory_order_relaxed) + bytes;
    allocation_counters.pmr_allocated_bytes_total.fetch_add(bytes, std::memory_order_relaxed);
    allocation_counters.pmr_allocation_count.fetch_add(1, std::memory_order_relaxed);
    update_peak(allocation_counters.pmr_peak_live_bytes, pmr_live);
}
void record_deallocation(void* pointer, bool pmr) {
    if (pointer == nullptr)
        return;
    const auto bytes = static_cast<std::uint64_t>(mi_usable_size(pointer));
    allocation_counters.live_bytes.fetch_sub(bytes, std::memory_order_relaxed);
    allocation_counters.deallocation_count.fetch_add(1, std::memory_order_relaxed);
    if (!pmr)
        return;
    allocation_counters.pmr_live_bytes.fetch_sub(bytes, std::memory_order_relaxed);
    allocation_counters.pmr_deallocation_count.fetch_add(1, std::memory_order_relaxed);
}
void free_tracked(void* pointer) noexcept {
    record_deallocation(pointer, false);
    mi_free(pointer);
}
class Mimalloc_Memory_Resource final : public std::pmr::memory_resource {
protected:
    void* do_allocate(std::size_t bytes, std::size_t alignment) override {
        auto pointer = mi_new_aligned(bytes, alignment);
        record_allocation(pointer, true);
        return pointer;
    }
    void do_deallocate(void* pointer, std::size_t, std::size_t) override {
        record_deallocation(pointer, true);
        mi_free(pointer);
    }
    bool do_is_equal(const std::pmr::memory_resource& other) const noexcept override {
        return this == &other;
    }
};
}
void* operator new(std::size_t bytes) {
    auto pointer = mi_new(bytes);
    record_allocation(pointer, false);
    return pointer;
}
void* operator new[](std::size_t bytes) {
    auto pointer = mi_new(bytes);
    record_allocation(pointer, false);
    return pointer;
}
void* operator new(std::size_t bytes, const std::nothrow_t&) noexcept {
    auto pointer = mi_new_nothrow(bytes);
    record_allocation(pointer, false);
    return pointer;
}
void* operator new[](std::size_t bytes, const std::nothrow_t&) noexcept {
    auto pointer = mi_new_nothrow(bytes);
    record_allocation(pointer, false);
    return pointer;
}
void operator delete(void* pointer) noexcept {
    free_tracked(pointer);
}
void operator delete[](void* pointer) noexcept {
    free_tracked(pointer);
}
void operator delete(void* pointer, const std::nothrow_t&) noexcept {
    free_tracked(pointer);
}
void operator delete[](void* pointer, const std::nothrow_t&) noexcept {
    free_tracked(pointer);
}
void operator delete(void* pointer, std::size_t) noexcept {
    free_tracked(pointer);
}
void operator delete[](void* pointer, std::size_t) noexcept {
    free_tracked(pointer);
}
void* operator new(std::size_t bytes, std::align_val_t alignment) {
    auto pointer = mi_new_aligned(bytes, static_cast<std::size_t>(alignment));
    record_allocation(pointer, false);
    return pointer;
}
void* operator new[](std::size_t bytes, std::align_val_t alignment) {
    auto pointer = mi_new_aligned(bytes, static_cast<std::size_t>(alignment));
    record_allocation(pointer, false);
    return pointer;
}
void* operator new(std::size_t bytes, std::align_val_t alignment, const std::nothrow_t&) noexcept {
    auto pointer = mi_new_aligned_nothrow(bytes, static_cast<std::size_t>(alignment));
    record_allocation(pointer, false);
    return pointer;
}
void* operator new[](std::size_t bytes, std::align_val_t alignment, const std::nothrow_t&) noexcept {
    auto pointer = mi_new_aligned_nothrow(bytes, static_cast<std::size_t>(alignment));
    record_allocation(pointer, false);
    return pointer;
}
void operator delete(void* pointer, std::align_val_t) noexcept {
    free_tracked(pointer);
}
void operator delete[](void* pointer, std::align_val_t) noexcept {
    free_tracked(pointer);
}
void operator delete(void* pointer, std::size_t, std::align_val_t) noexcept {
    free_tracked(pointer);
}
void operator delete[](void* pointer, std::size_t, std::align_val_t) noexcept {
    free_tracked(pointer);
}
void operator delete(void* pointer, std::align_val_t, const std::nothrow_t&) noexcept {
    free_tracked(pointer);
}
void operator delete[](void* pointer, std::align_val_t, const std::nothrow_t&) noexcept {
    free_tracked(pointer);
}
std::pmr::memory_resource* server_memory_resource() {
    static Mimalloc_Memory_Resource resource;
    return &resource;
}
void initialize_server_memory_resource() {
    mi_process_init();
    std::pmr::set_default_resource(server_memory_resource());
}
Server_Allocator_Snapshot server_allocator_snapshot() {
    Server_Allocator_Snapshot result;
    result.live_bytes = allocation_counters.live_bytes.load(std::memory_order_relaxed);
    result.peak_live_bytes = allocation_counters.peak_live_bytes.load(std::memory_order_relaxed);
    result.allocated_bytes_total = allocation_counters.allocated_bytes_total.load(std::memory_order_relaxed);
    result.allocation_count = allocation_counters.allocation_count.load(std::memory_order_relaxed);
    result.deallocation_count = allocation_counters.deallocation_count.load(std::memory_order_relaxed);
    result.pmr_live_bytes = allocation_counters.pmr_live_bytes.load(std::memory_order_relaxed);
    result.pmr_peak_live_bytes = allocation_counters.pmr_peak_live_bytes.load(std::memory_order_relaxed);
    result.pmr_allocated_bytes_total = allocation_counters.pmr_allocated_bytes_total.load(std::memory_order_relaxed);
    result.pmr_allocation_count = allocation_counters.pmr_allocation_count.load(std::memory_order_relaxed);
    result.pmr_deallocation_count = allocation_counters.pmr_deallocation_count.load(std::memory_order_relaxed);
    result.allocator_version = mi_version();
    return result;
}
void collect_server_memory(bool force) {
    mi_collect(force);
}