#include <stdlib.h>
#include <math.h>
#include <stdio.h>
#include <stdint.h>
#include "lisp.h"
#include <assert.h>

static Memory MEMORY;

static size_t alignToPow2(size_t n) {
  return pow(2, ceil(log2(n)));
}

static size_t alignToBlock(size_t n) {
  return alignToPow2(MAX(((n + sizeof(Block) - 1) / sizeof(Block)) + 1, 2));
}

void memory_init(size_t size) {
  MEMORY.used = 0;
  MEMORY.size = alignToPow2(MAX(size, 2));
  MEMORY.buffer = malloc(MEMORY.size * sizeof(Block));
  MEMORY.freelist = 0;
  MEMORY.buffer[MEMORY.freelist] = (Block) {
    .garbage = true,
    .size = MEMORY.size,
  };
  MEMORY.buffer[MEMORY.freelist].data->next = -1u;
}

void memory_free(void) {
  free(MEMORY.buffer);
}

Block* memory_get(Pointer data) {
  assert(data < MEMORY.size);
  return &MEMORY.buffer[data];
}

static void iterate_freelist(size_t free_size,
			     Pointer* free,
			     Block** free_previous,
			     bool free_break,
			     Pointer after_start,
			     size_t after_size,
			     Pointer* after,
			     Block** after_previous,
			     bool after_break,
			     Pointer* last,
			     Block** last_previous) {
  Pointer freelist = MEMORY.freelist;
  Block* previous = NULL;
  Block* block = NULL;
  while (freelist != -1u) {
    block = memory_get(freelist);
    if (block->size >= free_size) {
      if (free) *free = freelist;
      if (free_previous) *free_previous = previous;
      if (free_break) break;
    }

    if (freelist + block->size == MEMORY.size) {
      if (last) *last = freelist;
      if (last_previous) *last_previous = previous;
    }

    if (freelist == after_start &&
	block->size >= after_size) {
      if (after) *after = freelist;
      if (after_previous) *after_previous = previous;
      if (after_break) break;
    }

    previous = block;
    freelist = block->data->next;
  }
}

static Pointer use_free_slot(Type type, size_t size, Pointer free, Block* previous) {
  Block* block = memory_get(free);
  if (block->size == size) {
    if (previous != NULL) {
      previous->data->next = block->data->next;
    } else {
      MEMORY.freelist = block->data->next;
    }
    
    block->type = type;
    block->garbage = false;
    return free;
  } 
  
  size_t offset = size;
  size_t idx = free + offset;
  Block* new = block + offset;
  new->data->next = block->data->next;
  new->size = block->size - size;
  new->garbage = true;
  
  if (previous != NULL) {
    previous->data->next = idx;
  } else {
    MEMORY.freelist = idx;
  }
  
  block->type = type;
  block->garbage = false;
  block->size = size;

  return free;
}

Pointer alloc_free_slot(Type type, size_t size,  Pointer last, Block* last_previous) {
  // Memory_Get last free slot in buffer
  Pointer freelist = last;
  size_t requiredSize = MEMORY.size;
  if (last == -1u) {
    freelist = MEMORY.size;
    requiredSize += size;
  } else {
    if (last_previous != NULL) {
      last_previous->data->next = NEXT(last);
    } else {
      MEMORY.freelist = NEXT(last);
    }

    requiredSize += size - SIZE(last);
  }

  while (MEMORY.size < requiredSize) {
    MEMORY.size *= 2;
  }
  MEMORY.buffer = realloc(MEMORY.buffer, MEMORY.size * sizeof(Block));

  Block* block = memory_get(freelist);
  block->type = type;
  block->garbage = false;
  block->size = size;

  if (freelist + size < MEMORY.size) {
    block += size;
    block->size = MEMORY.size - freelist - size;
    block->garbage = true;
    block->data->next = MEMORY.freelist;
    MEMORY.freelist = freelist + size;
  } else {
    MEMORY.freelist = -1u;
  }

  return freelist;
}

Pointer memory_new(Type type, size_t size) {
  /* if (size == 0) { */
  /*   return -1u; // Why alloc an element of zero size? */
  /* } */

  size = alignToBlock(size);

  Pointer free = -1u;
  Block* free_previous = NULL;
  Pointer last = -1u;
  Block* last_previous = NULL;
  iterate_freelist(size,
		   &free,
		   &free_previous,
		   true,
		   0,
		   0,
		   NULL,
		   NULL,
		   false,
		   &last,
		   &last_previous);

  if (free != -1u) { 
    return use_free_slot(type, size, free, free_previous);
  }

  /* return alloc_free_slot(type, size, last, last_previous); */
  return alloc_free_slot(type, size, last, last_previous);
}

void memory_destroy(Pointer ptr) {
  Block* block = memory_get(ptr);
  Block* last = NULL;
  Block* beforeNext = NULL;

  Pointer freelist = MEMORY.freelist;
  Pointer previous = -1u;
  Pointer next = -1u;
  while (freelist != -1u &&
	 (previous == -1u || next == -1u)) {
    Block* free = memory_get(freelist);
    if (freelist + free->size == ptr) {
      previous = freelist;
    }

    if (freelist == ptr + block->size) {
      next = freelist;
      beforeNext = last;
    }
    last = free;
    freelist = free->data->next;
  }

  if (next != -1u) {
    Block* nblk = memory_get(next);
    block->size += nblk->size;
    if (beforeNext != NULL) {
      last->data->next = nblk->data->next;
    } else {
      MEMORY.freelist = nblk->data->next;
    }
  }

  if (previous != -1u) {
    Block* pblk = memory_get(previous);
    pblk->size += block->size;
  } else {
    block->data->next = MEMORY.freelist;
    MEMORY.freelist = ptr;
  }
}

Pointer memory_resize(Pointer ptr, size_t new_size) {
  Block* pblock = memory_get(ptr);
  Block* block;
  size_t prev_size = pblock->size;
  new_size = MAX(alignToBlock(new_size), 2);

  size_t remaining = new_size - prev_size;

  if (new_size == prev_size) {
    return ptr;
  }

  if (new_size < prev_size) {
    // Shrink
    Pointer rest = ptr + new_size;
    block = memory_get(rest);
    block->size = prev_size - new_size;
    block->data->next = MEMORY.freelist;
    block->garbage = true;
    MEMORY.freelist = rest;
    pblock->size = new_size;
    return ptr;
  }

  Pointer free = -1u;
  Block* free_previous = NULL;
  Pointer after = -1u;
  Block* after_previous = NULL;
  Pointer last = -1u;
  Block* last_previous = NULL;
  iterate_freelist(new_size,
		   &free,
		   &free_previous,
		   false,
		   ptr + prev_size,
		   remaining,
		   &after,
		   &after_previous,
		   true,
		   &last,
		   &last_previous);

  if (after != -1u) {
    block = memory_get(after);
    if (block->size > remaining) {
      Pointer restPtr = after + remaining;
      Block* rest = memory_get(restPtr);
      rest->size = block->size - remaining;
      NEXT(restPtr) = NEXT(after);
      if (after_previous != NULL) {
	after_previous->data->next = restPtr;
      } else {
	MEMORY.freelist = restPtr;
      }
    } else {
      if (after_previous != NULL) {
	after_previous->data->next = block->data->next;
      } else {
	MEMORY.freelist = block->data->next;
      }
    }
    block->garbage = false;
    pblock->size = new_size;
    return ptr;
  }


  if (free != -1u ||
      (ptr + pblock->size != last &&
       ptr + pblock->size != MEMORY.size)) {
    if (free == -1u) {
      free = alloc_free_slot(pblock->type, new_size, last, last_previous);
    } else {
      free = use_free_slot(pblock->type, new_size, free, free_previous);
    }

    Pointer offset = free - ptr;
    for (size_t i = 0; i < MEMORY.size; ) {
      block = memory_get(i);
      switch (block->type) {
      case CONS:
	if (CONS(i).car == ptr) CONS(i).car += offset;
	if (CONS(i).cdr == ptr) CONS(i).cdr += offset;
	break;
      case TABLE: // TODO: handle for Table
      case ARRAY: // TODO: handle for Vector
      default: break;
      }
      i += block->size;
    }
    
    // Copy old memory to new memory
    for (size_t i = 1; i <= prev_size; i++) {
      MEMORY.buffer[free + i] = MEMORY.buffer[ptr + i];
    }
    
    // Free old memory
    memory_destroy(ptr);
    
    return free;
  }

  alloc_free_slot(pblock->type, remaining, last, last_previous);
  return ptr;
}