#include "vulkan.h"

#include <inttypes.h>
#include <stdbool.h>
#include <stdlib.h>
#include <string.h>
#include <vulkan/vulkan_core.h>

#define GPU_MAX_ALLOCS 1024
#define MAX_PAGES 10
#define MIN_PAGE_SIZE 262144

struct GpuAlloc {
  uint64_t start;
  uint64_t size;
  struct GpuAlloc *next;
};

struct GpuPage {
  uint64_t size;
  uint count;
  uint32_t index;
  VkDeviceMemory memory;
  struct GpuAlloc allocs[GPU_MAX_ALLOCS];
  struct GpuAlloc *alloc;
  void *mapped;
};

struct GpuMem {
  struct GpuPage pages[MAX_PAGES + 1];
};

struct GpuMem *make_gpu_alloc() {
  struct GpuMem *mem = malloc(sizeof(struct GpuMem));
  memset(mem, 0, sizeof(struct GpuMem));

  return mem;
}

void free_gpu_page(struct Vk *state, struct GpuPage *page) {
  vkFreeMemory(state->device, page->memory, NULL);
  memset(page, 0, sizeof(struct GpuPage));
}

void free_gpu_mem(struct Vk *state, struct GpuMem *mem) {
  for (int i = 0; i < MAX_PAGES; i++) {
    if (*(uint32_t *)&mem->pages[i] == 0) {
      break;
    }
    free_gpu_page(state, &mem->pages[i]);
  }
  memset(mem, 0, sizeof(struct GpuMem));
  free(mem);
}

void insert_gpu_page(struct Vk *state, struct GpuMem *mem, uint32_t index,
                     uint64_t size, bool mapped) {
  int i = 0;
  while (*(int *)&mem->pages[i] != 0) {
    i++;
  }
  if (i == MAX_PAGES) {
    crash("out of gpu memory");
  }

  struct GpuPage *page = &mem->pages[i];

  VkMemoryAllocateInfo alloc_info = {0};
  alloc_info.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
  alloc_info.allocationSize = size;
  alloc_info.memoryTypeIndex = index;
  CHECK_VK_RESULT(
      vkAllocateMemory(state->device, &alloc_info, NULL, &page->memory));
  if (mapped) {
    vkMapMemory(state->device, page->memory, 0, size, 0, &page->mapped);
    meow("page %d is safe to use up till %p", i, page->mapped + size);
  } else {
    page->mapped = NULL;
  }
  page->index = index;
  page->size = size;
  page->allocs[0] = (struct GpuAlloc){.start = size, .size = 0, .next = NULL};
  page->allocs[1] =
      (struct GpuAlloc){.start = 0, .size = 0, .next = &page->allocs[0]};
  page->alloc = &page->allocs[1];
}

static struct GpuAlloc *get_next_gpu_alloc(struct GpuPage *page) {
  while (page->allocs[page->count].start != 0 ||
         page->allocs[page->count].next != NULL) {
    page->count++;
    if (page->count >= GPU_MAX_ALLOCS) {
      crash("out of allocs!");
    }
  }
  return &page->allocs[page->count];
}

struct GpuPointer gpu_mem_malloc(struct Vk *state, struct GpuMem *mem,
                                 uint64_t size, uint32_t index, bool mapped) {
  size += 64 - (size % 64);
  for (int i = 0; i < MAX_PAGES; i++) {
    if (*(uint32_t *)&mem->pages[i] == 0) {
      insert_gpu_page(state, mem, index,
                      size < MIN_PAGE_SIZE ? MIN_PAGE_SIZE : size, mapped);
    }
    if (mem->pages[i].index == index) {
      struct GpuPage *page = &mem->pages[i];
      struct GpuAlloc *prev = page->alloc;
      while (prev->next != NULL) {
        struct GpuAlloc *next = prev->next;
        if (next->start - (prev->start + prev->size) >= size) {
          struct GpuAlloc *new = get_next_gpu_alloc(page);
          *new = (struct GpuAlloc){
              prev->start + prev->size,
              size,
              next,
          };
          prev->next = new;
          meow("GPUMALLOC %d:%ld offset %lu", i, new - page->allocs,
               new->start);
          return (struct GpuPointer){new->start, page->memory,
                                     page->mapped + new->start};
        } else {
          prev = next;
        }
      }
    }
  }
  crash("oops no more gpu pages");
  return (struct GpuPointer){0, NULL, NULL};
}

void gpu_mem_free(struct Vk *state, struct GpuMem *mem,
                  struct GpuPointer pointer) {
  for (int i = 0; i < MAX_PAGES; i++) {
    if (mem->pages[i].memory == pointer.memory) {
      struct GpuPage *page = &mem->pages[i];
      struct GpuAlloc *current = page->alloc;
      struct GpuAlloc *prev = page->alloc;
      while (current != NULL &&
             (pointer.offset > current->start || current->size == 0)) {
        prev = current;
        current = current->next;
      }
      if (current == NULL || pointer.offset != current->start) {
        meow("WHOOOPPSIIIEEEE could not find allocated block, potential dobble "
             "free");
        paw_print_backtrace();
        return;
      }
      prev->next = current->next;
      int index = current - page->allocs;
      if (index < page->count) {
        page->count = index;
      }
      memset(&page->allocs[index], 0, sizeof(struct GpuAlloc));
      meow("GPUFREE %d:%ld", i, current - page->allocs);
      return;
    }
  }
  crash("page not tracked! pointer mem is %p", pointer.memory);
}