#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;

#define CANARY 0x5a

// 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 {
  union {
    uint32_t header;
    struct {
      /* Is this memory block currently in use (hasn't been kfree'd) */
      uint8_t used : 1;
      /* Number of words allocated. Does not include the header. */
      uint16_t size : 12;
      /* The location of the previous block (in WORDS from offest) */
      kalloc_off_t prev : 12;
      uint8_t canary : 7;
    } PACKED;
  };

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

static_assert(offsetof(kalloc_node_t, mem) == 4, "Offset check failed.");

kalloc_node_t* kalloc_start;

#define kalloc_node_out_of_range(node) \
  ((uint8_t*)(node) == ((uint8_t*)&DATA_SEGMENT_STOP) + MAX_HEAP_SIZE)

#define kalloc_node_next(cur) \
  ((kalloc_node_t*)(((uint8_t*)(cur)) + (((cur)->size + 1) * 4)))

#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)*4))

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

#define get_kalloc_node(mem) ((kalloc_node_t*)(((uint8_t*)mem) - 4))

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

void* kalloc(size_t size)
{
  if (!kalloc_start) {
    kalloc_start = (kalloc_node_t*)DATA_SEGMENT_STOP_ADDR;
    memset(kalloc_start, 0, sizeof(kalloc_node_t));
    kalloc_start->size = (MAX_HEAP_SIZE / 4) - 1;
    kalloc_start->canary = CANARY;
  }

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

    if (!cur->used && (cur->size >= realsz)) {
      cur->used = true;
      size_t orig_size = cur->size;
      cur->size = realsz;

      if (orig_size > realsz) {
        /* This kalloc node needs to split into two blocks. */
        kalloc_node_t* next = kalloc_node_next(cur);
        next->used = 0;
        next->size = orig_size - realsz - sizeof(kalloc_node_t) / 4;
        next->prev = offset;
        next->canary = CANARY;

        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_t) / 4) + cur->size;
  }

  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 (!cur->used && cur != kalloc_start) {
    cur = kalloc_node_prev(cur);
  }

  if (cur == kalloc_start && !cur->used) {
    last_freed = cur;
  } else {
    last_freed = kalloc_node_next(cur);
  }

  /* Find the next used block. */
  cur = orig;
  while (!cur->used && !kalloc_node_out_of_range(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 = ((uint8_t*)next_used - (last_freed->mem)) / 4;
}

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 (!header->used) {
    panic("Heap double free or corruption!\n");
    return;
  }

  header->used = 0;
  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");
  kalloc_node_t* cur = kalloc_node_at_off(0);

  while (!kalloc_node_out_of_range(cur)) {
    printf(
        "header (%04x) {used=%d, size=%5d, prev=%04x, canary=%02x}\n",
        kalloc_node_get_off(cur),
        cur->used,
        cur->size,
        cur->prev,
        cur->canary);
    cur = kalloc_node_next(cur);
  }
}

/* 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 (cur->canary != CANARY) {
      snprintf(
          error,
          len,
          "Node has corrupted canary. %02x vs expected %02x\n",
          cur->canary,
          CANARY);
      return 1;
    }

    total_size += cur->size + 1;

    kalloc_node_t* next = kalloc_node_next(cur);
    if ((uint8_t*)next == ((uint8_t*)&DATA_SEGMENT_STOP) + MAX_HEAP_SIZE) {
      break;
    } else if (
        (uint8_t*)next > (uint8_t*)DATA_SEGMENT_STOP_ADDR + MAX_HEAP_SIZE) {
      snprintf(
          error,
          len,
          "Next node points is out of bounds. %p vs max of %p\n",
          next,
          (void*)(DATA_SEGMENT_STOP_ADDR + MAX_HEAP_SIZE));
      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 * 4,
        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;
}

#endif