#include <knob/format.h>

#include <pland/syscall.h>

#include <stdbool.h>
#include <stdint.h>

//structure appears at the start of each block
struct block_header {
  struct block_header *next;
  struct block_header *prev;

  //does not include header
  uint32_t length;
  bool allocated;
};

static struct block_header *first_block = 0;

static void add_header(struct block_header *bh) {
  bh->next = first_block;
  bh->prev = 0;
  if (first_block)
    first_block->prev = bh;
  first_block = bh;
}

__attribute__ ((malloc))
void *get_block(uint32_t bytes) {
  struct block_header *header = 0;
  for (struct block_header *ptr = first_block; ptr; ptr = ptr->next) {
    if (ptr->allocated)
      continue;
    if (ptr->length == bytes) {
      header = ptr;
      break;
    }
    if (ptr->length > bytes + sizeof(struct block_header)) {
      struct block_header *bh = (void *)ptr + sizeof(struct block_header) + bytes;

      bh->length = ptr->length - sizeof(struct block_header) - bytes;
      bh->allocated = false;
      add_header(bh);

      ptr->length = bytes;
      header = ptr;
      break;
    }
  }

  if (!header) {
    uint32_t size_with_header = bytes + sizeof(struct block_header);
    if (!(size_with_header % 4096)) {
      header = _allocate_ram(size_with_header / 4096);
      if (!header)
        return 0;
      header->length = bytes;
      add_header(header);
    }
    else {
      uint32_t pages = (bytes + sizeof(struct block_header) * 2) / 4096 + 1;
      header = _allocate_ram(pages);
      if (!header)
        return 0;
      header->length = bytes;
      add_header(header);

      struct block_header *bh = (void *)header + sizeof(struct block_header) + bytes;
      bh->length = pages * 4096 - bytes - 2 * sizeof(struct block_header);
      bh->allocated = false;
      add_header(bh);
    }
  }

  header->allocated = true;
  return (void *)header + sizeof(struct block_header);
}

static void remove_header(struct block_header *bh) {
  if (bh == first_block)
    first_block = bh->next;
  if (bh->prev)
    bh->prev->next = bh->next;
  if (bh->next)
    bh->next->prev = bh->prev;
}

void free_block(void *block) {
  struct block_header *header = block - sizeof(struct block_header);

  header->allocated = false;

  void *after = block + header->length;
  for (struct block_header *ptr = first_block; ptr; ptr = ptr->next) {
    if (ptr == after) {
      if (!ptr->allocated) {
        header->length += ptr->length + sizeof(struct block_header);
        remove_header(ptr);
      }
      break;
    }
  }

  //we traverse the linked list twice. it would probably be more efficient
  //to do both finding the after and finding the before in the same loop.
  for (struct block_header *ptr = first_block; ptr; ptr = ptr->next)
    if ((void *)ptr + sizeof(struct block_header) + ptr->length == header) {
      if (!ptr->allocated) {
        ptr->length += sizeof(struct block_header) + header->length;
        remove_header(header);
      }
      break;
    }
}