//
// Created by Dan Crosby on 4/6/24.
//

#include "FlatbufferWrapper.hpp"

#include <sys/mman.h>
#include <unistd.h>

namespace shardok {

// Uses mmap and munmap to directly allocate pages
class MmapAllocator final : public WrapperAllocator {
public:
    auto allocate(const size_t size) -> uint8_t* override {
        auto* b = static_cast<uint8_t*>(
                mmap(nullptr, size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0));

        if (b == MAP_FAILED) {
            printf("mmap failed\n");
            abort();
        }
        return b;
    }
    auto deallocate(uint8_t* buf, const size_t size) -> void override {
        if (const auto result = munmap(buf, size); result != 0) {
            printf("munmap failed, errno %s\n", strerror(errno));
            abort();
        }
    }

    ~MmapAllocator() override = default;
};

// Uses malloc and free to allocate memory
class MallocAllocator final : public WrapperAllocator {
public:
    auto allocate(const size_t size) -> uint8_t* override {
        auto* b = static_cast<uint8_t*>(malloc(size));

        if (b == nullptr) {
            printf("malloc failed");
            abort();
        }
        return b;
    }
    auto deallocate(uint8_t* buf, const size_t size) -> void override { free(buf); }

    ~MallocAllocator() override = default;
};

// Uses a stack to reuse memory allocated by one of the other allocators
class ReusingAllocator final : public WrapperAllocator {
private:
    std::unique_ptr<WrapperAllocator> fallbackAllocator;

    const size_t pageSize = getpagesize();

    auto RoundUpToPage(const size_t size) const -> size_t {
        return (size + pageSize - 1) & ~(pageSize - 1);
    }

    struct StackItem {
        uint8_t* buffer;
        size_t size;
        StackItem* next;
    };

    StackItem* usedHead;
    StackItem* emptyHead;

    std::mutex mtx;

public:
    explicit ReusingAllocator(std::unique_ptr<WrapperAllocator> alloc)
        : fallbackAllocator(std::move(alloc)),
          usedHead(nullptr),
          emptyHead(nullptr) {}

    auto allocate(size_t size) -> uint8_t* override {
        uint8_t* buffer = nullptr;
        size = RoundUpToPage(size);

        {
            std::unique_lock lk(mtx);

            while (buffer == nullptr) {
                if (auto* top = usedHead; top != nullptr) {
                    if (top->size >= size) {
                        buffer = top->buffer;
                        usedHead = top->next;
                        top->next = emptyHead;
                        emptyHead = top;
                    } else {
                        fallbackAllocator->deallocate(top->buffer, top->size);
                        usedHead = top->next;
                        top->next = emptyHead;
                        emptyHead = top;
                    }
                } else {
                    buffer = fallbackAllocator->allocate(size);
                }
            }
        }

        return buffer;
    }
    auto deallocate(uint8_t* buf, size_t size) -> void override {
        size = RoundUpToPage(size);

        std::unique_lock lk(mtx);
        auto* newTop = emptyHead;
        if (newTop) {
            emptyHead = newTop->next;
        } else {
            newTop = new StackItem();
        }
        newTop->buffer = buf;
        newTop->size = size;
        newTop->next = usedHead;
        usedHead = newTop;
    }

    ~ReusingAllocator() override {
        std::unique_lock lk(mtx);
        while (const auto* top = usedHead) {
            usedHead = top->next;
            fallbackAllocator->deallocate(top->buffer, top->size);
            delete top;
        }
        while (const auto* top = emptyHead) {
            emptyHead = top->next;
            delete top;
        }
    }
};

auto WrapperAllocator::GetDefaultAllocator() -> WrapperAllocator* {
    static pthread_once_t once = PTHREAD_ONCE_INIT;
    static WrapperAllocator* allocator = nullptr;
    pthread_once(&once, []() {
        // Change this to use one of the other allocators
        allocator = new MallocAllocator();
    });

    return allocator;
}

}  // namespace shardok