#include "Memory.h"
#include <array>
#include <atomic>
#include <memory>
#include <mutex>
#include <new>
#include <stdexcept>
namespace renderive {
namespace {
constexpr std::size_t memory_domain_count = static_cast<std::size_t>(Memory_Domain::Other) + 1;

std::size_t memory_domain_index(Memory_Domain domain) noexcept {
    std::size_t index = static_cast<std::size_t>(domain);
    return index < memory_domain_count ? index : static_cast<std::size_t>(Memory_Domain::Other);
}

void* default_allocate(void*, std::size_t size, std::size_t alignment) {
    return std::pmr::new_delete_resource()->allocate(size, alignment);
}
void default_deallocate(void*, void* address, std::size_t size, std::size_t alignment) {
    std::pmr::new_delete_resource()->deallocate(address, size, alignment);
}
class Callback_Memory_Resource final : public std::pmr::memory_resource {
public:
    Callback_Memory_Resource() : callbacks{nullptr, default_allocate, default_deallocate} {}
    void set_callbacks(const Memory_Callbacks& value) {
        callbacks = value.allocate ? value : Memory_Callbacks{nullptr, default_allocate, default_deallocate};
    }
    Memory_Stats stats() const noexcept {
        Memory_Stats result;
        result.upstream_allocation_count = allocation_count.load(std::memory_order_relaxed);
        result.upstream_deallocation_count = deallocation_count.load(std::memory_order_relaxed);
        result.upstream_current_bytes = current_bytes.load(std::memory_order_relaxed);
        result.upstream_peak_bytes = peak_bytes.load(std::memory_order_relaxed);
        return result;
    }
private:
    void* do_allocate(std::size_t size, std::size_t alignment) override {
        void* result = callbacks.allocate(callbacks.context, size, alignment);
        if (!result)
            throw std::bad_alloc();
        allocation_count.fetch_add(1, std::memory_order_relaxed);
        std::uint64_t current = current_bytes.fetch_add(size, std::memory_order_relaxed) + size;
        std::uint64_t peak = peak_bytes.load(std::memory_order_relaxed);
        while (peak < current && !peak_bytes.compare_exchange_weak(peak, current, std::memory_order_relaxed)) {}
        return result;
    }
    void do_deallocate(void* address, std::size_t size, std::size_t alignment) override {
        callbacks.deallocate(callbacks.context, address, size, alignment);
        deallocation_count.fetch_add(1, std::memory_order_relaxed);
        current_bytes.fetch_sub(size, std::memory_order_relaxed);
    }
    bool do_is_equal(const std::pmr::memory_resource& other) const noexcept override {
        return this == &other;
    }
    Memory_Callbacks callbacks;
    std::atomic<std::uint64_t> allocation_count{};
    std::atomic<std::uint64_t> deallocation_count{};
    std::atomic<std::uint64_t> current_bytes{};
    std::atomic<std::uint64_t> peak_bytes{};
};
class alignas(64) Domain_Memory_Resource final : public std::pmr::memory_resource {
public:
    void set_domain(Memory_Domain value) noexcept {
        domain = value;
    }
    Memory_Domain_Stats stats() const noexcept {
        Memory_Domain_Stats result;
#if RENDERIVE_MEMORY_DOMAIN_STATS
        result.allocation_count = allocation_count.load(std::memory_order_relaxed);
        result.deallocation_count = deallocation_count.load(std::memory_order_relaxed);
        result.current_bytes = current_bytes.load(std::memory_order_relaxed);
        result.peak_bytes = peak_bytes.load(std::memory_order_relaxed);
#endif
        return result;
    }
private:
    void* do_allocate(std::size_t size, std::size_t alignment) override;
    void do_deallocate(void* address, std::size_t size, std::size_t alignment) override;
    bool do_is_equal(const std::pmr::memory_resource& other) const noexcept override {
        return this == &other;
    }
    Memory_Domain domain = Memory_Domain::Other;
#if RENDERIVE_MEMORY_DOMAIN_STATS
    std::atomic<std::uint64_t> allocation_count{};
    std::atomic<std::uint64_t> deallocation_count{};
    std::atomic<std::uint64_t> current_bytes{};
    std::atomic<std::uint64_t> peak_bytes{};
#endif
};
struct Memory_System {
    Memory_System() {
        for (std::size_t i = 0; i < domains.size(); ++i)
            domains[i].set_domain(static_cast<Memory_Domain>(i));
    }
    std::mutex mutex;
    Callback_Memory_Resource upstream;
    std::unique_ptr<std::pmr::synchronized_pool_resource> small_pool;
    std::atomic<std::pmr::memory_resource*> small_resource{};
    std::atomic_bool initialized{};
    bool configured{};
    std::array<Domain_Memory_Resource, memory_domain_count> domains;
};
Memory_System& memory_system() {
    static Memory_System value;
    return value;
}
bool domain_uses_small_pool(Memory_Domain domain) {
    switch (domain) {
    case Memory_Domain::Renderable:
    case Memory_Domain::Scheduler_Task:
    case Memory_Domain::Update_Completion:
    case Memory_Domain::Other:
        return true;
    default:
        return false;
    }
}
std::pmr::memory_resource* upstream_resource(Memory_System& system) {
    system.initialized.store(true, std::memory_order_release);
    return &system.upstream;
}
std::pmr::memory_resource* small_object_resource(Memory_System& system) {
    system.initialized.store(true, std::memory_order_release);
#if RENDERIVE_MEMORY_GLOBAL_POOL
    std::pmr::memory_resource* resource = system.small_resource.load(std::memory_order_acquire);
    if (resource)
        return resource;
    std::lock_guard<std::mutex> lock(system.mutex);
    resource = system.small_resource.load(std::memory_order_relaxed);
    if (resource)
        return resource;
    std::pmr::pool_options options;
    options.max_blocks_per_chunk = 64;
    options.largest_required_pool_block = 4 * 1024;
    system.small_pool = std::make_unique<std::pmr::synchronized_pool_resource>(options, &system.upstream);
    resource = system.small_pool.get();
    system.small_resource.store(resource, std::memory_order_release);
    return resource;
#else
    return &system.upstream;
#endif
}
std::pmr::memory_resource* domain_resource(Memory_Domain domain) {
    Memory_System& system = memory_system();
    if (domain_uses_small_pool(domain))
        return small_object_resource(system);
    return upstream_resource(system);
}
void* Domain_Memory_Resource::do_allocate(std::size_t size, std::size_t alignment) {
    void* result = domain_resource(domain)->allocate(size, alignment);
#if RENDERIVE_MEMORY_DOMAIN_STATS
    allocation_count.fetch_add(1, std::memory_order_relaxed);
    std::uint64_t current = current_bytes.fetch_add(size, std::memory_order_relaxed) + size;
    std::uint64_t peak = peak_bytes.load(std::memory_order_relaxed);
    while (peak < current && !peak_bytes.compare_exchange_weak(peak, current, std::memory_order_relaxed)) {}
#endif
    return result;
}
void Domain_Memory_Resource::do_deallocate(void* address, std::size_t size, std::size_t alignment) {
    domain_resource(domain)->deallocate(address, size, alignment);
#if RENDERIVE_MEMORY_DOMAIN_STATS
    deallocation_count.fetch_add(1, std::memory_order_relaxed);
    current_bytes.fetch_sub(size, std::memory_order_relaxed);
#endif
}
}
void set_memory_callbacks(const Memory_Callbacks& callbacks) {
    if (!!callbacks.allocate != !!callbacks.deallocate)
        throw std::invalid_argument("allocate and deallocate must be configured together");
    Memory_System& system = memory_system();
    std::lock_guard<std::mutex> lock(system.mutex);
    if (system.configured || system.initialized.load(std::memory_order_acquire))
        throw std::logic_error("Core memory resource is already configured");
    system.upstream.set_callbacks(callbacks);
    system.configured = true;
}
Memory_Stats memory_stats() noexcept {
    return memory_system().upstream.stats();
}
Memory_Domain_Stats memory_domain_stats(Memory_Domain domain) noexcept {
    return memory_system().domains[memory_domain_index(domain)].stats();
}
std::pmr::memory_resource* memory_resource() {
    return small_object_resource(memory_system());
}
std::pmr::memory_resource* memory_resource(Memory_Domain domain) {
#if RENDERIVE_MEMORY_DOMAIN_STATS
    Memory_System& system = memory_system();
    system.initialized.store(true, std::memory_order_release);
    return &system.domains[memory_domain_index(domain)];
#else
    return domain_resource(domain);
#endif
}
void release_unused_memory() {
    Memory_System& system = memory_system();
    std::lock_guard<std::mutex> lock(system.mutex);
    if (system.small_pool)
        system.small_pool->release();
}
void shutdown_memory() {
    Memory_System& system = memory_system();
    std::lock_guard<std::mutex> lock(system.mutex);
    system.small_resource.store(nullptr, std::memory_order_release);
    system.small_pool.reset();
    system.initialized.store(false, std::memory_order_release);
    system.configured = false;
    system.upstream.set_callbacks(Memory_Callbacks{});
}
}