// Same dependencies as vulkan memory allocator
#include "../../l1/core/uint_segments.h"
#include "../../l1/core/util.h"
#include "../../l1_5/core/buff_rb_tree_node.h"
#include "../../../gen/l1_5/BufRBTree_MapU64ToU64.h"

typedef struct MargaretBufAllocatorOneBlock MargaretBufAllocatorOneBlock;

typedef struct {
    MargaretBufAllocatorOneBlock* block;
    U64 start;
    U64 len;
} MargaretBAFreeSegment;

bool MargaretBAFreeSegment_less_len(const MargaretBAFreeSegment* A, const MargaretBAFreeSegment* B){
    if (A->len == B->len) {
        if (A->block == B->block) {
            return A->start < B->start;
        }
        return (uintptr_t)A->block < (uintptr_t)B->block;
    }
    return A->len < B->len;
}

U64 margaret_bump_buffer_size_to_alignment(U64 A, U8 alignment_exp){
    if (A & ((1ull << alignment_exp) - 1))
        A = A - (A & ((1ull << alignment_exp) - 1)) + (1ull << alignment_exp);
    return A;
}

typedef struct {
    MargaretBufAllocatorOneBlock* block;
    U64 start;
    U64 len;
} MargaretSubbuf;

struct MargaretBufAllocatorOneBlock{
    BufRBTree_MapU64ToU64 occupants;
    U64 capacity;
    U64 occupation_counter;
    VkDeviceMemory mem_hand;
    VkBuffer buf_hand;
    void* mapped_memory;
};

void MargaretBufAllocatorOneBlock_drop(MargaretBufAllocatorOneBlock self){
    BufRBTree_MapU64ToU64_drop(self.occupants);
}
#include "../../../gen/l1/eve/margaret/ListMargaretBufAllocatorOneBlock.h"

#include "../../../gen/l1/eve/margaret/OptionMargaretBAFreeSegment.h"
#include "../../../gen/l1/eve/margaret/VecMargaretBAFreeSegment.h"
#include "../../../gen/l1_5/eve/margaret/BufRBTreeByLen_SetMargaretBAFreeSegment.h"

typedef struct {
    ListMargaretBufAllocatorOneBlock blocks;
    BufRBTreeByLen_SetMargaretBAFreeSegment mem_free_space;
    VkDevice device;
    VkPhysicalDevice physical_device;
    VkBufferUsageFlags usage;
    U8 memory_type_id;
    U8 alignment_exp;
    bool host_visible;
} MargaretBufAllocator;


void MargaretBufAllocator__erase_gap(
    MargaretBufAllocator* self, MargaretBufAllocatorOneBlock* block, U64 start, U64 len){
    if (len == 0)
        return;
    bool eret = BufRBTreeByLen_SetMargaretBAFreeSegment_erase(&self->mem_free_space,
        &(MargaretBAFreeSegment){.block = block, .start = start, .len = len});
    assert(eret);
    block->occupation_counter += len;
    assert(block->occupation_counter <= block->capacity);
}

void MargaretBufAllocator__insert_gap(
    MargaretBufAllocator* self, MargaretBufAllocatorOneBlock* block, U64 start, U64 len){
    if (len == 0)
        return;
    bool iret = BufRBTreeByLen_SetMargaretBAFreeSegment_insert(&self->mem_free_space,
        (MargaretBAFreeSegment){.block = block, .start = start, .len = len});
    assert(iret);
    assert(len <= block->occupation_counter);
    block->occupation_counter -= len;
}

OptionMargaretBAFreeSegment MargaretBufAllocator__search_gap(MargaretBufAllocator* self, U64 req_size){
    assert(req_size % (1ull << self->alignment_exp) == 0);
    U64 sit = BufRBTreeByLen_SetMargaretBAFreeSegment_find_min_grtr_or_eq(&self->mem_free_space,
        &(MargaretBAFreeSegment){.len = req_size});
    if (sit == 0)
        return None_MargaretBAFreeSegment();
    return Some_MargaretBAFreeSegment(*BufRBTreeByLen_SetMargaretBAFreeSegment_at_iter(&self->mem_free_space, sit));
}

void MargaretBufAllocator__add_block(MargaretBufAllocator* self, U64 capacity){
    VkBuffer buffer;
    check(vkCreateBuffer(self->device, &(VkBufferCreateInfo){
        .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
        .size = capacity,
        .usage = self->usage,
        .sharingMode = VK_SHARING_MODE_EXCLUSIVE
    }, NULL, &buffer) == VK_SUCCESS);
    VkMemoryRequirements memory_requirements;
    vkGetBufferMemoryRequirements(self->device, buffer, &memory_requirements);
    VkDeviceMemory memory;
    check(vkAllocateMemory(self->device, &(VkMemoryAllocateInfo){
        .sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO,
        .allocationSize = memory_requirements.size,
        .memoryTypeIndex = self->memory_type_id
    }, NULL, &memory) == VK_SUCCESS);
    check(vkBindBufferMemory(self->device, buffer, memory, 0) == VK_SUCCESS);
    void* mapped_memory = NULL;
    if (self->host_visible) {
        check(vkMapMemory(self->device, memory, 0, capacity, 0, &mapped_memory) == VK_SUCCESS);
    }
    ListMargaretBufAllocatorOneBlock_insert(&self->blocks, (MargaretBufAllocatorOneBlock){
        .occupants = BufRBTree_MapU64ToU64_new_reserved(1),
        .capacity = capacity,
        .occupation_counter = capacity,
        .mem_hand = memory, .buf_hand = buffer, .mapped_memory = mapped_memory
        });
}

MargaretBufAllocator MargaretBufAllocator_new(
    VkDevice device, VkPhysicalDevice physical_device,
    VkBufferUsageFlags usage, U8 memory_type_id, U8 alignment_exp, bool host_visible, U64 initial_block_size
    ){
    MargaretBufAllocator self = {
        .blocks = ListMargaretBufAllocatorOneBlock_new(),
        .mem_free_space = BufRBTreeByLen_SetMargaretBAFreeSegment_new_reserved(1),
        .device = device, .physical_device = physical_device, .usage = usage, .memory_type_id = memory_type_id,
        .alignment_exp = alignment_exp, .host_visible = host_visible
    };
    MargaretBufAllocator__add_block(&self, initial_block_size);
    MargaretBufAllocator__insert_gap(&self, &self.blocks.first->el, 0, initial_block_size);
    return self;
}

void MargaretBufAllocator__put_buf_to_a_gap(MargaretBufAllocator* self, MargaretBAFreeSegment segment, U64 req_size){
    assert(req_size <= segment.len);
    MargaretBufAllocator__erase_gap(self, segment.block, segment.start, segment.len);
    MargaretBufAllocator__insert_gap(self, segment.block,
        segment.start + req_size, segment.len - req_size);
    BufRBTree_MapU64ToU64* occupants = &segment.block->occupants;
    bool iret = BufRBTree_MapU64ToU64_insert(occupants, segment.start, req_size);
    assert(iret);
}

U64Segment MargaretBufAllocator__get_left_free_space(
    const MargaretBufAllocator* self, const MargaretSubbuf* allocation){
    U64 occ_start = allocation->start;
    U64 prev_occ_it = BufRBTree_MapU64ToU64_find_max_less(&allocation->block->occupants, allocation->start);
    if (prev_occ_it != 0) {
        U64 prev_occ_start;
        U64 prev_occ_taken_size;
        BufRBTree_MapU64ToU64_at_iter(&allocation->block->occupants, prev_occ_it, &prev_occ_start, &prev_occ_taken_size);

        assert(prev_occ_start + prev_occ_taken_size <= occ_start);
        return (U64Segment){
            .start = prev_occ_start + prev_occ_taken_size,
            .len = occ_start - (prev_occ_start + prev_occ_taken_size)};
    }
    return (U64Segment){.start = 0, .len = occ_start};
}

U64Segment MargaretBufAllocator__get_right_free_space(
    const MargaretBufAllocator* self, const MargaretSubbuf* allocation){
    U64 occ_start = allocation->start;
    U64 occ_taken_size = allocation->len;

    U64 next_occ_it = BufRBTree_MapU64ToU64_find_min_grtr(&allocation->block->occupants, allocation->start);
    if (next_occ_it != 0) {
        U64 next_occ_start;
        U64 next_occ_taken_size;
        BufRBTree_MapU64ToU64_at_iter(&allocation->block->occupants, next_occ_it, &next_occ_start, &next_occ_taken_size);
        assert(occ_start + occ_taken_size <= next_occ_start);
        return (U64Segment){.start = occ_start + occ_taken_size, .len = next_occ_start - (occ_start + occ_taken_size)};
    }
    return (U64Segment){.start = occ_start + occ_taken_size, .len = allocation->block->capacity - (occ_start + occ_taken_size)};
}


void MargaretBufAllocator_drop(MargaretBufAllocator self){
    for (ListNodeMargaretBufAllocatorOneBlock* bi = self.blocks.first; bi; bi = bi->next) {
        vkDestroyBuffer(self.device, bi->el.buf_hand, NULL);
        vkFreeMemory(self.device, bi->el.mem_hand, NULL);
    }
    ListMargaretBufAllocatorOneBlock_drop(self.blocks);
    BufRBTreeByLen_SetMargaretBAFreeSegment_drop(self.mem_free_space);
}

/* Idk how to hide this monster */
void MargaretBufAllocator_debug(const MargaretBufAllocator* self){
    if (!self->host_visible)
        return;
    printf(" ======== MargaretBufAllocator state ======== \n");
    int n_segments = (int)self->mem_free_space.el.len;
    printf("Blocks:\n");
    for (ListNodeMargaretBufAllocatorOneBlock* block_it = self->blocks.first; block_it; block_it = block_it->next) {
        U64 free_space_acc_segs = 0;
        U64 occ_space_acc_occ = 0;
        MargaretBufAllocatorOneBlock* block = &block_it->el;
        int n_occupants = (int)block->occupants.el.len;
        printf("-*- occupied: %lu/%lu, occupants: %d\n", block->occupation_counter, block->capacity, n_occupants);
        for (int si = 0; si < n_segments; si++) {
            MargaretBAFreeSegment fseg = self->mem_free_space.el.buf[si];
            if (fseg.block == block) {
                assert(fseg.start + fseg.len <= block->capacity);
                free_space_acc_segs += fseg.len;
            }
        }
        for (int oi = 0; oi < n_occupants; oi++) {
            KVPU64ToU64 occ = block->occupants.el.buf[oi];
            assert(occ.key + occ.value <= block->capacity);
            occ_space_acc_occ += occ.value;
            for (int sc = 0; sc < n_occupants; sc++) {
                KVPU64ToU64 occ2 = block->occupants.el.buf[sc];
                if (sc != oi) {
                    assert(occ.key + occ.value <= occ2.key || occ2.key + occ2.value <= occ.key);
                }
            }
        }
        assert(free_space_acc_segs == block->capacity - block->occupation_counter);
        assert(occ_space_acc_occ == block->occupation_counter);
    }
}

/* Free one subbuffer, not a whole MBA :)  */
void MargaretBufAllocator_free(MargaretBufAllocator* self, MargaretSubbuf allocation){
    U64Segment left_free_space = MargaretBufAllocator__get_left_free_space(self, &allocation);
    U64Segment right_free_space = MargaretBufAllocator__get_right_free_space(self, &allocation);

    MargaretBufAllocator__erase_gap(self, allocation.block, left_free_space.start, left_free_space.len);
    MargaretBufAllocator__erase_gap(self, allocation.block, right_free_space.start, right_free_space.len);
    MargaretBufAllocator__insert_gap(self, allocation.block,
        left_free_space.start,
        right_free_space.start + right_free_space.len - left_free_space.start);

    bool eret = BufRBTree_MapU64ToU64_erase(&allocation.block->occupants, allocation.start);
    assert(eret);
    // MargaretBufAllocator_debug(self);
}

NODISCARD MargaretSubbuf MargaretBufAllocator_alloc(MargaretBufAllocator* self, U64 req_size){
    // MargaretBufAllocator_debug(self);
    req_size = margaret_bump_buffer_size_to_alignment(req_size, self->alignment_exp);

    VkPhysicalDeviceMaintenance3Properties maintenance3_properties = {
        .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_MAINTENANCE_3_PROPERTIES,
    };
    VkPhysicalDeviceProperties2 properties = {
        .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2,
        .pNext = &maintenance3_properties,
    };
    vkGetPhysicalDeviceProperties2(self->physical_device, &properties);

    check(req_size <= maintenance3_properties.maxMemoryAllocationSize);
    OptionMargaretBAFreeSegment free_gap = MargaretBufAllocator__search_gap(self, req_size);
    if (free_gap.variant == Option_None) {
        assert(self->blocks.first != NULL);
        U64 pitch = self->blocks.first->el.capacity;
        // Old blocks remain intact
        U64 new_capacity = MAX_U64(req_size, MIN_U64(2 * pitch, maintenance3_properties.maxMemoryAllocationSize));
        MargaretBufAllocator__add_block(self, new_capacity);
        MargaretBufAllocatorOneBlock* new_block = &self->blocks.first->el;
        MargaretBufAllocator__insert_gap(self, new_block, req_size, new_capacity - req_size);
        new_block->occupation_counter = req_size;
        bool iret = BufRBTree_MapU64ToU64_insert(&new_block->occupants, 0, req_size);
        assert(iret);
        // MargaretBufAllocator_debug(self);
        return (MargaretSubbuf){.block = &self->blocks.first->el, 0, req_size};
    }
    MargaretBufAllocator__put_buf_to_a_gap(self, free_gap.some, req_size);
    // MargaretBufAllocator_debug(self);
    return (MargaretSubbuf){.block = free_gap.some.block, .start = free_gap.some.start, req_size};
}

void MargaretBufAllocator_shrink(MargaretBufAllocator* self, MargaretSubbuf* allocation, U64 smaller_size){
    smaller_size = margaret_bump_buffer_size_to_alignment(smaller_size, self->alignment_exp);
    assert(smaller_size > 0);
    assert(smaller_size <= allocation->len);

    U64Segment right_free_space = MargaretBufAllocator__get_right_free_space(self, allocation);
    MargaretBufAllocator__erase_gap(self, allocation->block, right_free_space.start, right_free_space.len);
    MargaretBufAllocator__insert_gap(self, allocation->block,
        allocation->start + smaller_size,
        right_free_space.len + (allocation->len - smaller_size));

    allocation->len = smaller_size;
}

/* It actually may returns a 'null-MargaretBuf-allocation' : if return value .len field is zero it means
 * that expansion in-place was possible and the allocator argument was updated with a new size and nothing was returned.
 * But if ret value .len field is non-zero it means a valid MargaretSubbuf object was returned and the
 * `allocation` argument was untouched. It remains a valid object, you need to deallocate it yourself
 */
NODISCARD MargaretSubbuf MargaretBufAllocator_expand(
    MargaretBufAllocator* self, MargaretSubbuf* allocation, U64 bigger_size){
    bigger_size = margaret_bump_buffer_size_to_alignment(bigger_size, self->alignment_exp);

    U64Segment right_free_space = MargaretBufAllocator__get_right_free_space(self, allocation);
    if (allocation->start + bigger_size > right_free_space.start + right_free_space.len){
        return MargaretBufAllocator_alloc(self, bigger_size);
    }
    // MargaretBufAllocator_debug(self);
    MargaretBufAllocator__erase_gap(self, allocation->block, right_free_space.start, right_free_space.len);
    MargaretBufAllocator__insert_gap(self, allocation->block,
        allocation->start + bigger_size,
        right_free_space.len + (allocation->len - bigger_size));

    allocation->len = bigger_size;
    U64 my_it = BufRBTree_MapU64ToU64_find(&allocation->block->occupants, allocation->start);
    assert(my_it > 0 && my_it <= allocation->block->occupants.el.len);
    allocation->block->occupants.el.buf[my_it - 1].value = bigger_size;
    // MargaretBufAllocator_debug(self);
    return (MargaretSubbuf){0};
}

char* MargaretSubbuf_get_mapped(const MargaretSubbuf* allocation){
    assert(allocation->block->mapped_memory);
    assert(allocation->start + allocation->len <= allocation->block->capacity);
    return (char*)allocation->block->mapped_memory + allocation->start;
}

VkBuffer MargaretSubbuf_get_buffer(const MargaretSubbuf* allocation){
    assert(allocation->start + allocation->len <= allocation->block->capacity);
    return allocation->block->buf_hand;
}

/* It tries to expand buffer, but if it fails, it creates a freshly-new buffer. It copies all
 * the data from old buffer to new one and frees the old buffer, while replacing
 * info in `allocation` variable with info about new allocation.
 */
void MargaretBufAllocator_expand_or_move_old_host_visible(
    MargaretBufAllocator* self, MargaretSubbuf* allocation, U64 bigger_size){
    assert(self->host_visible);
    MargaretSubbuf maybe_bigger = MargaretBufAllocator_expand(self, allocation, bigger_size);
    if (maybe_bigger.len > 0) {
        memcpy(MargaretSubbuf_get_mapped(&maybe_bigger), MargaretSubbuf_get_mapped(allocation), allocation->len);
        MargaretBufAllocator_free(self, *allocation);
        *allocation = maybe_bigger;
        // MargaretBufAllocator_debug(self);
    }
    // MargaretBufAllocator_debug(self);
}

/* It tries to expand buffer, but if it fails, it creates a freshly-new buffer. It
 * frees the old buffer, while replacing
 * info in `allocation` variable with info about new allocation. Old data gets lost
 */
void MargaretBufAllocator_expand_or_free_old(
    MargaretBufAllocator* self, MargaretSubbuf* allocation, U64 bigger_size){
    MargaretSubbuf maybe_bigger = MargaretBufAllocator_expand(self, allocation, bigger_size);
    if (maybe_bigger.len > 0) {
        MargaretBufAllocator_free(self, *allocation);
        *allocation = maybe_bigger;
    }
}