#include "arch.h"
#include "kern/mem.h"
#include "kern/common.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 halloc_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 HALLOC_NODE {
  union {
    uint32_t header;
    struct {
      /* Is this memory block currently in use (hasn't been hfree'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) */
      halloc_off_t prev:12;
      uint8_t canary:7;
    } PACKED;
  };

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

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

halloc_node_t* halloc_start;

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

#define halloc_node_next(cur) \
  ((halloc_node_t*)(((uint8_t*)(cur)) + (((cur)->size + 1) * 4)))

#define halloc_node_prev(cur) halloc_node_at_off(cur->prev)

#define halloc_node_at_off(offset) \
  ((halloc_node_t*)(((uint8_t*) halloc_start) + (offset) * 4))

#define halloc_node_get_off(node) \
  (((uint32_t)(((uint8_t*)(node)) - ((uint8_t*)(halloc_start)))) / 4)

#define get_halloc_node(mem) \
  ((halloc_node_t*)(((uint8_t*)mem) - 4))

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

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

  size_t realsz = size_for(size); /* Clip the size to the nearest word. */
  halloc_off_t offset = 0; 
  while (offset < (MAX_HEAP_SIZE / 4)) {
    halloc_node_t* cur = halloc_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 halloc node needs to split into two blocks. */
        halloc_node_t* next = halloc_node_next(cur);
        next->used = 0;
        next->size = orig_size - realsz - sizeof(halloc_node_t) / 4;
        next->prev = offset;
        next->canary = CANARY;

        halloc_node_t* nextnext = halloc_node_next(next);
        if (halloc_node_get_off(nextnext) < (MAX_HEAP_SIZE / 4)) {
          nextnext->prev = halloc_node_get_off(next);
        }
      }

      return (void*) cur->mem;
    }

    offset += (sizeof(halloc_node_t) / 4) + cur->size;
  }

  return NULL;
}


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

  /* Find the earliest contiguous free'd block. */
  while (!cur->used && cur != halloc_start) {
    cur = halloc_node_prev(cur);
  }

  if (cur == halloc_start && !cur->used) {
    last_freed = cur;
  } else {
    last_freed = halloc_node_next(cur);
  }

  /* Find the next used block. */
  cur = orig;
  while (!cur->used && !halloc_node_out_of_range(cur)) {
    cur = halloc_node_next(cur);
  }

  next_used = cur;

  if (!halloc_node_out_of_range(next_used)) {
    next_used->prev = halloc_node_get_off(last_freed);
  }

  last_freed->size = ((uint8_t*) next_used - (last_freed->mem)) / 4;
}

#ifdef FOR_TESTING
#include <assert.h>
#include <stdio.h>
void panic(const char* x)
{
  fprintf(stderr, "%s\n", x);
  assert(0);
}
#else
void panic(const char* x)
{
  for(;;);
}
#endif

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

  halloc_node_t* header = get_halloc_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_halloc_get_next_ptr(void* ptr)
{
  halloc_node_t* node = ptr - sizeof(halloc_node_t);
  halloc_node_t* next = halloc_node_next(node);
  
  return next->mem;
}

void* debug_halloc_get_prev_ptr(void* ptr)
{
  halloc_node_t* node = ptr - sizeof(halloc_node_t);
  halloc_node_t* prev = halloc_node_prev(node);
  
  return prev->mem;
}

void debug_print_blocks()
{
  printf("------ Print Blocks -------\n");
  halloc_node_t* cur = halloc_node_at_off(0);

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

/* Tests that we can walk up and down the allocated blocks and that they
 * are properly aligned. */
int debug_halloc_assert_consistency(char* error, size_t len)
{
  halloc_node_t* cur = halloc_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;

    halloc_node_t* next = halloc_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 == halloc_start) {
    return 0;
  }

  while (loop_check < 10000) {
    halloc_node_t* prev = halloc_node_prev(cur);
    ++ n_blocks_back;

    if (prev == halloc_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