#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 notrun1 = 0;

void *MALLOCCHERI(size_t sz)
{
   if (notrun1 == 0){
     INITREGULARALLOC();
     notrun1 = 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 += 1;
//   return MALLOCCHERI(sz);
// }

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