#include <errno.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>

#include <stdlib.h>
#include	<sys/types.h>

#include <cheriintrin.h>
#include <cheri/cheric.h>

#include <sys/stat.h>
#include <fcntl.h>

#include <assert.h>
#include <sys/time.h>
#include <sys/errno.h>
#include <stdint.h>
#include <stdio.h>
#include	<unistd.h>

#define MAXPAGESIZES    2

static char *heap_start;
static char *heap;
static size_t HEAP_SIZE = 1024 * 1024 * 1024;

void *ptr;
int MallocCounter;
int malloc_called = 0;

size_t sizeUsed;

// Instrcutor allocator to create the huge page
// of 1 GB 
// __attribute__((constructor))
static void INITREGULARALLOC() {
   size_t sz;
   // Hard-coded for 1GB huge page
   sz = 1073741824;

   int error, fd, pscnt, pn;

   size_t ps[2];

   size_t size[3];

  //  pn = getpagesizes(size, 3);
  //  printf("page size is [%d]", size[2]);

  //  pscnt = pagesizes(ps);

	fd = shm_create_largepage(SHM_ANON, O_CREAT | O_RDWR, 1, SHM_LARGEPAGE_ALLOC_DEFAULT, 0);

   if (fd < 0 && errno == ENOTTY) {
      perror("sh_create_largepages");
      close(fd);
      exit(EXIT_FAILURE);
   }
   // if (fd < 0)
   //     perror("no large page supported");
   //     exit(EXIT_FAILURE);
	// if (fd < 0 && errno == ENOTTY)
	// 	atf_tc_skip("no large page support");
	// ATF_REQUIRE_MSG(fd >= 0, "shm_create_largepage failed; errno=%d", errno);

	if (ftruncate(fd, sz) < 0) {
        perror("ftruncate");
        close(fd);
        exit(EXIT_FAILURE);
    }
	// if (error != 0 && errno == ENOMEM)
	// 	/*
	// 	 * The test system might not have enough memory to accommodate
	// 	 * the request.
	// 	 */
	// 	atf_tc_skip("failed to allocate %zu-byte superpage", sz);
	// ATF_REQUIRE_MSG(error == 0, "ftruncate failed; errno=%d", errno);

   ptr = mmap(NULL, sz,
    PROT_READ|PROT_WRITE, MAP_SHARED,fd,0);

   // Added error handling
    if(ptr == MAP_FAILED)
    {
        perror("mmap");
        exit(EXIT_FAILURE);
    }


  //  fprintf(stderr, "heap used alloc %lu\n", heap - heap_start);

    MallocCounter = (int)sz;
}

// -- Custom malloc and free functions written
// This will be replaced with mmap since we already
// do a initial mmap. 

int notrun = 0;

void *MALLOCCHERI(size_t sz)
{
   if (notrun == 0){
     INITREGULARALLOC();
     notrun = 1;
   }

   sz = __builtin_align_up(sz, _Alignof(max_align_t));

   // printf("%d \n", sz);
  //  printf("%d Malloc counter\n", MallocCounter);

   MallocCounter -= sz;
   void *ptrLink = &ptr[MallocCounter];
   ptrLink = cheri_setbounds(ptrLink, sz);

   return ptrLink;

//   if (heap + sz > heap_start + HEAP_SIZE) return NULL;
//   heap += sz;
//   return heap - sz;

}

// Quick cheri free implementation
void FREECHERI(void *ptr) {

  //  printf("free called \n");

   // get bounds from
   int len = cheri_getlen(ptr);

  //  printf("free len %d \n", len);

   munmap(ptr, len);
}

__attribute__((destructor))
static void malloc_exit() {
  fprintf(stderr, "heap used %lu\n", malloc_called);
}

// void *malloc(size_t sz) {
//   // if (!heap) heap = heap_start = mmap(NULL, HEAP_SIZE,
//   //   PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANON,-1,0);

//   // char *new_ptr = __builtin_align_up(
//   //   heap, -cheri_representable_alignment_mask(sz));
//   // size_t bounds = cheri_representable_length(sz);
//   // sz = __builtin_align_up(sz, _Alignof(max_align_t));

//   // if (new_ptr + sz > heap_start + HEAP_SIZE)
//   //   return NULL;
//   // heap = new_ptr + sz;
//   // return cheri_bounds_set_exact(new_ptr, bounds);
//   return MALLOCCHERI(sz);
// }

void *malloc(size_t sz) {
  // if (!heap) heap = heap_start = mmap(NULL, HEAP_SIZE,
  //   PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANON,-1,0);
  // sz = __builtin_align_up(sz, _Alignof(max_align_t));
  // if (heap + sz > heap_start + HEAP_SIZE) return NULL;
  // heap += sz;
  // return heap - sz;
  malloc_called += 2;
  return MALLOCCHERI(sz);
}

void free(void *ptr) { 
  FREECHERI(ptr);
}

void *realloc(void *ptr, size_t sz) {
  void *new_ptr = malloc(sz);
  if (ptr && new_ptr) memmove(new_ptr, ptr, sz);
  return new_ptr;
}

void *calloc(size_t nmemb, size_t sz) {
  char *ptr = malloc(nmemb * sz);
  bzero(ptr, nmemb * sz);
  return ptr;
}

void *reallocarray(void *ptr, size_t nmemb, size_t sz) {
  return realloc(ptr, nmemb * sz);
}

void *recallocarray(void *ptr, size_t oldnmemb, size_t nmemb, size_t sz) {
  void *new_ptr = malloc(nmemb * sz);
  if (ptr && new_ptr) memmove(new_ptr, ptr, oldnmemb * sz);
  if (new_ptr && nmemb > oldnmemb) bzero(new_ptr + oldnmemb * sz, (nmemb - oldnmemb) * sz);
  return new_ptr;
}

void freezero(void *ptr, size_t sz) { }

void *aligned_alloc(size_t alignment, size_t sz) {
  return malloc(sz);
}

void *malloc_conceal(size_t sz) {
  return malloc(sz);
}

void *calloc_conceal(size_t nmemb, size_t sz) {
  return calloc(nmemb, sz);
}

int posix_memalign(void **ptr, size_t alignment, size_t sz) {
  *ptr = malloc(sz);
  return *ptr == 0 ? 0 : ENOMEM;
}

void *memalign(size_t alignment, size_t sz) {
  return malloc(sz);
}

void *valloc(size_t sz) {
  return malloc(sz);
}