#include "Frame_Memory.h"
#include "Memory.h"
#include <algorithm>
#include <memory>
#include <vector>
namespace renderive {
namespace {
constexpr std::size_t Frame_Arena_Min_Size = 64 * 1024;
constexpr std::size_t Frame_Arena_Max_Retained_Size = 64 * 1024 * 1024;
}
struct Thread_Frame_Scratch;
class Frame_Overflow_Resource final : public std::pmr::memory_resource {
public:
    explicit Frame_Overflow_Resource(Thread_Frame_Scratch* scratch) : scratch(scratch) {}
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;
    }
    Thread_Frame_Scratch* scratch{};
};
struct Thread_Frame_Scratch {
    Thread_Frame_Scratch() {
        storage.resize(Frame_Arena_Min_Size);
        desired_size = Frame_Arena_Min_Size;
    }
    std::pmr::memory_resource* overflow_upstream() {
        if (!overflow_pool)
            overflow_pool = std::make_unique<std::pmr::unsynchronized_pool_resource>(memory_resource(Memory_Domain::Frame_Arena));
        return overflow_pool.get();
    }
    void begin_frame(std::size_t initial_buffer_size) {
        frame_overflow_count = 0;
        frame_overflow_bytes = 0;
        std::size_t target_size = std::max(initial_buffer_size, desired_size);
        target_size = std::max(target_size, Frame_Arena_Min_Size);
        target_size = std::min(target_size, Frame_Arena_Max_Retained_Size);
        if (storage.size() < target_size) {
            storage.resize(target_size);
        }
    }
    void finish_frame() {
        std::uint64_t frame_bytes = static_cast<std::uint64_t>(storage.size() + frame_overflow_bytes);
#if RENDERIVE_FRAME_ARENA_GROW
        if (frame_overflow_bytes) {
            std::size_t next_size = static_cast<std::size_t>(frame_bytes + frame_bytes / 8 + 4096);
            desired_size = std::clamp(next_size, Frame_Arena_Min_Size, Frame_Arena_Max_Retained_Size);
        }
#endif
        frame_overflow_count = 0;
        frame_overflow_bytes = 0;
    }
    std::vector<std::byte> storage;
    std::unique_ptr<std::pmr::unsynchronized_pool_resource> overflow_pool;
    Frame_Overflow_Resource overflow_resource{this};
    std::size_t desired_size{};
    std::size_t frame_overflow_count{};
    std::size_t frame_overflow_bytes{};
    int depth{};
};
void* Frame_Overflow_Resource::do_allocate(std::size_t size, std::size_t alignment) {
    scratch->frame_overflow_count += 1;
    scratch->frame_overflow_bytes += size;
    return scratch->overflow_upstream()->allocate(size, alignment);
}
void Frame_Overflow_Resource::do_deallocate(void* address, std::size_t size, std::size_t alignment) {
    scratch->overflow_upstream()->deallocate(address, size, alignment);
}
namespace {
thread_local std::pmr::memory_resource* current_frame_resource{};
Thread_Frame_Scratch& thread_frame_scratch() {
    thread_local Thread_Frame_Scratch value;
    return value;
}
}
Frame_Arena::Frame_Arena(std::size_t initial_buffer_size) {
    scratch = &thread_frame_scratch();
    previous = current_frame_resource;
    if (scratch->depth == 0) {
        scratch->begin_frame(initial_buffer_size);
        arena.emplace(scratch->storage.data(), scratch->storage.size(), &scratch->overflow_resource);
    }
    else {
        arena.emplace(&scratch->overflow_resource);
    }
    ++scratch->depth;
    current_frame_resource = resource();
}
Frame_Arena::~Frame_Arena() {
    current_frame_resource = previous;
    bool outer = scratch && scratch->depth == 1;
    arena.reset();
    if (outer)
        scratch->finish_frame();
    --scratch->depth;
}
std::pmr::memory_resource* Frame_Arena::resource() noexcept {
    return arena ? &*arena : memory_resource();
}
std::pmr::memory_resource* frame_memory_resource() {
    return current_frame_resource ? current_frame_resource : memory_resource();
}
}