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

// 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) */
      bool     used:1;
      /* Number of words allocated. Does not include the header. */
      uint16_t size:15;
      /* The location of the previous block (in WORDS from offest) */
      halloc_off_t prev;
    } PACKED;
  };

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

halloc_node_t* halloc_start;

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

  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;

        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)
{
  if (cur->used) return;

  halloc_node_t* next = halloc_node_next(cur);
  if (!next->used) {
    cur->size += next->size + sizeof(halloc_node_t) / 4;
  }

  if (cur != halloc_start) {
    coalesce(halloc_node_prev(cur));
  }
}

void hfree(void* mem)
{
  /* TODO this should be a critical section. */
  if (!mem) {
    /* Consistent with C, freeing a NULL pointer does nothing. */
    return;
  }

  halloc_node_t* header = ((halloc_node_t*) mem) - 1;

  if (!header->used) {
#ifdef FOR_TESTING
    assert(header->used);
#endif
    return; // TODO make this fail on ARM.
  }

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

/* 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;

  while(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;
  }

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

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

    if (prev == halloc_start)  {
      return 0;
    }

    cur = prev;
    ++ loop_check;
  }

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

#endif