#include "kern/mem.h"

#include "arch.h"
#include "kern/common.h"
#include "kern/panic.h"

#ifdef ARCH_STM32L4
/* Provide a definition for memset() when not provided for the
 * microcontroller. */
void* memset(void* dest, int c, size_t n)
{
  uint8_t c8 = (uint8_t)c;
  uint8_t* dest8 = (uint8_t*)dest;
  uint8_t* to = dest8 + n;

  while (dest8 < to) {
    *(dest8++) = c8;
  }

  return dest;
}

#else

void* memset(void* dest, int c, size_t n);

#endif

typedef uint16_t kalloc_off_t;

// The sizes will count the number of WORDS allocated.
// Since there's a max size of 16k, only 12 bits will be
// needed for this.
typedef struct KALLOC_NODE {
  uint16_t size_words;
  kalloc_off_t prev;
  /*
   * LSB is whether this block is used.
   *
   * Rest of the bits must equal (0xdeadbeee >> 1)
   */
  uint32_t used_and_canary;

  uint8_t mem[]; /* The memory to use. */
} PACKED kalloc_node_t;

#ifdef ARCH_PC
typedef uint64_t ptrint_t;
#else
typedef uint32_t ptrint_t;
#endif

#define CANARY ((uint32_t)0xdeadbeee)
#define kalloc_node_in_use(node) ((node)->used_and_canary & 1)
#define kalloc_node_get_canary(node) ((node)->used_and_canary & (~1))
#define WORD_SIZE (sizeof(uint32_t))
#define SIZEOF_KALLOC_NODE_WORDS (sizeof(kalloc_node_t) / WORD_SIZE)
#define HEAP_START_ADDR ((ptrint_t)&HEAP_START)
#define REAL_HEAP_START *(uint8_t*)(HEAP_START_ADDR + (4 - HEAP_START_ADDR % 4))
#define MAX_HEAP_SIZE ((&HEAP_STOP - &REAL_HEAP_START))
#define MAX_HEAP_SIZE_WORDS (MAX_HEAP_SIZE / WORD_SIZE)
#define kalloc_node_out_of_range(node) ((void*)(node) >= (void*)&HEAP_STOP)

#define kalloc_node_next(cur) \
  ((kalloc_node_t*)(((uint8_t*)(cur)) + (((cur)->size_words + SIZEOF_KALLOC_NODE_WORDS) * WORD_SIZE)))

#define kalloc_node_prev(cur) kalloc_node_at_off(cur->prev)

#define kalloc_node_at_off(offset) \
  ((kalloc_node_t*)(((uint8_t*)kalloc_start) + (offset)*WORD_SIZE))

#define kalloc_node_get_off(node) \
  ((uint32_t)(((((uint8_t*)(node)) - ((uint8_t*)(kalloc_start)))) / WORD_SIZE))

#define get_kalloc_node(mem) \
  ((kalloc_node_t*)(((uint8_t*)mem) - SIZEOF_KALLOC_NODE_WORDS * WORD_SIZE))

#define size_for(n) (((n) / 4) + ((n) % 4 != 0))

kalloc_node_t* kalloc_start;

void kalloc_init()
{
  kalloc_start = (kalloc_node_t*)&REAL_HEAP_START;
  memset(kalloc_start, 0, sizeof(kalloc_node_t));
  kalloc_start->size_words = MAX_HEAP_SIZE_WORDS - SIZEOF_KALLOC_NODE_WORDS;
  kalloc_start->used_and_canary = CANARY;
}

void* kalloc(size_t size)
{
  if (!kalloc_start) {
    kalloc_init();
  }

  size_t realsz_words = size_for(size); /* Clip the size to the nearest word. */
  kalloc_off_t offset = 0;
  while (offset < (MAX_HEAP_SIZE_WORDS)) {
    kalloc_node_t* cur = kalloc_node_at_off(offset);

    if (!kalloc_node_in_use(cur) && (cur->size_words >= realsz_words)) {
      cur->used_and_canary |= 1;
      size_t orig_size_words = cur->size_words;

      if (orig_size_words < realsz_words + SIZEOF_KALLOC_NODE_WORDS * 2) {
        /* If the original size is only slightly larger than the size to
         * allocate, there's no point in saving the leftover space. */
        realsz_words = orig_size_words;
      }

      cur->size_words = realsz_words;

      if (orig_size_words > realsz_words) {
        /* This kalloc node needs to split into two blocks. */
        kalloc_node_t* next = kalloc_node_next(cur);
        next->used_and_canary = CANARY;
        next->size_words =
            orig_size_words - realsz_words - SIZEOF_KALLOC_NODE_WORDS;
        next->prev = offset;

        kalloc_node_t* nextnext = kalloc_node_next(next);
        if (kalloc_node_get_off(nextnext) < (MAX_HEAP_SIZE / 4)) {
          nextnext->prev = kalloc_node_get_off(next);
        }
      }

      return (void*)cur->mem;
    }

    offset += SIZEOF_KALLOC_NODE_WORDS + cur->size_words;
  }

  return NULL;
}

/* Joins this node with the previous and next nodes if they're free. */
static void coalesce(kalloc_node_t* cur)
{
  kalloc_node_t* orig = cur;
  kalloc_node_t* last_freed;
  kalloc_node_t* next_used;

  /* Find the earliest contiguous free'd block. */
  while (!kalloc_node_in_use(cur) && cur != kalloc_start) {
    cur = kalloc_node_prev(cur);
  }

  if (cur == kalloc_start && !kalloc_node_in_use(cur)) {
    last_freed = cur;
  } else {
    last_freed = kalloc_node_next(cur);
  }

  /* Find the next used block. */
  cur = orig;
  while (!kalloc_node_out_of_range(cur) && !kalloc_node_in_use(cur)) {
    cur = kalloc_node_next(cur);
  }

  next_used = cur;

  if (!kalloc_node_out_of_range(next_used)) {
    next_used->prev = kalloc_node_get_off(last_freed);
  }

  last_freed->size_words =
      ((uint8_t*)next_used - (last_freed->mem)) / WORD_SIZE;
}

void kfree(void* mem)
{
  /* Like normal free(), do nothing on free'ing NULL */
  if (!mem) return;

  kalloc_node_t* header = get_kalloc_node(mem);
  if (!kalloc_node_in_use(header)) {
    panic("Heap double free or corruption!\n");
    return;
  }

  header->used_and_canary &= ~1;
  coalesce(header);
}

#ifdef FOR_TESTING

#include <stdio.h>

void* debug_kalloc_get_next_ptr(void* ptr)
{
  kalloc_node_t* node = ptr - sizeof(kalloc_node_t);
  kalloc_node_t* next = kalloc_node_next(node);

  return next->mem;
}

void* debug_kalloc_get_prev_ptr(void* ptr)
{
  kalloc_node_t* node = ptr - sizeof(kalloc_node_t);
  kalloc_node_t* prev = kalloc_node_prev(node);

  return prev->mem;
}

void debug_print_blocks()
{
  printf("------ Print Blocks -------\n");
  printf("MAX_HEAP_SIZE      : %ld\n", MAX_HEAP_SIZE);
  printf("MAX_HEAP_SIZE_WORDS: %ld\n", MAX_HEAP_SIZE_WORDS);
  printf("HEAP_START         : %p\n", &HEAP_START);
  printf("HEAP_STOP          : %p\n", &HEAP_STOP);
  printf("---------------------------\n");
  kalloc_node_t* cur = kalloc_node_at_off(0);

  int total_words = 0;
  int total_blocks = 0;
  while (!kalloc_node_out_of_range(cur)) {
    printf(
        "header (%04x)@%p {used=%d, size=%5d, prev=%04x, canary=%04x}\n",
        kalloc_node_get_off(cur),
        cur->mem,
        kalloc_node_in_use(cur),
        cur->size_words,
        cur->prev,
        kalloc_node_get_canary(cur));
    total_words += cur->size_words;
    total_blocks ++;
    cur = kalloc_node_next(cur);
  }

  printf("Total words allocated: %d\n", total_words);
  printf("Total blocks allocated: %d\n", total_blocks);
}

/* Tests that we can walk up and down the allocated blocks and that they
 * are properly aligned. */
int debug_kalloc_assert_consistency(char* error, size_t len)
{
  kalloc_node_t* cur = kalloc_node_at_off(0);
  size_t total_size = 0;
  size_t loop_check = 0;
  size_t n_blocks = 1;
  size_t n_blocks_back = 1;

  while (1) {
    if (kalloc_node_get_canary(cur) != CANARY) {
      snprintf(
          error,
          len,
          "Node has corrupted canary. %02x vs expected %02x\n",
          kalloc_node_get_canary(cur),
          CANARY);
      return 1;
    }

    total_size += cur->size_words + SIZEOF_KALLOC_NODE_WORDS;

    kalloc_node_t* next = kalloc_node_next(cur);
    if ((void*)next == (void*)&HEAP_STOP) {
      break;
    } else if ((void*)next > (void*)&HEAP_STOP) {
      snprintf(
          error,
          len,
          "Next node points is out of bounds. %p vs max of %p\n",
          next,
          &HEAP_STOP);
      return 1;
    }

    cur = next;
    ++n_blocks;
  }

  if (total_size * 4 != MAX_HEAP_SIZE) {
    snprintf(
        error,
        len,
        "Total recorded size is inconsistent. %lu vs %lu\n",
        total_size * WORD_SIZE,
        MAX_HEAP_SIZE);
    return 1;
  }

  if (cur == kalloc_start) {
    return 0;
  }

  while (loop_check < 10000) {
    kalloc_node_t* prev = kalloc_node_prev(cur);
    ++n_blocks_back;

    if (prev == kalloc_start) {
      if (n_blocks != n_blocks_back) {
        snprintf(
            error,
            len,
            "Different number of blocks found on the way back. Found %lu on "
            "the way back vs %lu up.\n",
            n_blocks_back,
            n_blocks);
        return 1;
      }
      return 0;
    }

    cur = prev;
    ++loop_check;
  }

  snprintf(error, len, "Loop check failed.\n");
  return 1;
}

int debug_is_heap_empty()
{
  return (void*)((uint8_t*)kalloc_start + kalloc_start->size_words * sizeof(uint32_t) + sizeof(kalloc_node_t)) ==
         (void*)&HEAP_STOP;
}

#endif