#include <stdlib.h>
#include <stdio.h>
#include <pthread.h>
#include <stdint.h>
#include "tlb_malloc.h"
#include "list_operations.h"
//16 * 4 * 4096
#define MEMORY_4K_PAGE 262144
#define PAGE_SIZE_4K 4096
#define PAGE_SIZE_2M 2097152
#define BYTES_FOR_ALLOCATION_LENGTH 4
#define ASSOC_MULTIPLIER_4K 1
/* Global variables */

///////////////////////   COLOR MALLOC

#include "./arch/arm/include/asm/cacheparams.h"
#include "./include/litmus/rt_param.h"

#include "./include/color.h"

#define PAGES_TO_ALLOC  5

static void* color_malloc_or_exit(size_t size ,void *virtual_address)
{
        void *mem;
        mem = color_malloc(size,virtual_address);
        if (!mem) {
                fprintf(stderr, "could not allocate memory.\n");
                exit(EXIT_FAILURE);
        }
        return mem;
}

static void check_memory(const unsigned char *start, const size_t len, const unsigned char expect)
{
        const unsigned char *cur = start;
        int i;

        for (i = 0; i < len; i++, cur++) {
                if (expect != *cur) {
                        fprintf(stderr, "inconsistent memory: "
                                        "start:0x%p  i:%d  cur:0x%p  "
                                        "expect:0x%x  have:0x%x\n",
                                        start, i, cur,
                                        expect, *cur);
			  exit(EXIT_FAILURE);
                }
        }
}

#define NR_COLORS_FILE "/proc/sys/litmus/color/cache/nr_colors"

void setup_pages(struct color_ctrl_page *ctrl, color_t start,
                color_t nr_pages, color_t stride)
{
        const int nr_colors = read_int_from_file(NR_COLORS_FILE);
        color_t color = start;
        int i;

        if (1 > nr_colors) {
                fprintf(stderr, "bad nr_colors\n");
                exit(EXIT_FAILURE);
        }

        for (i = 0; i < nr_pages; i++) {
                ctrl->colors[i] = color;
                ctrl->pages[i] = 1;
                color = (color + stride) % nr_colors;
        }
}


//////////////////////////////////////////
//variables to keep information about the DTLB 
__thread unsigned volatile int no_of_dtlb_sets_4k;
__thread unsigned volatile int dtlb_assoc_4k;
__thread unsigned volatile int no_of_dtlb_sets_for_array_4k = 8; 
__thread unsigned volatile int no_of_dtlb_sets_for_non_array_4k;
__thread unsigned volatile int no_of_dtlb_sets_2m;
__thread unsigned volatile int dtlb_assoc_2m;
__thread unsigned volatile int no_of_dtlb_sets_for_array_2m; 
__thread unsigned volatile int no_of_dtlb_sets_for_non_array_2m; 
__thread unsigned volatile int assoc_multiplier;
//keeps the address of the huge address space allocated
__thread void * memory = NULL;
__thread unsigned int memory_req_4k_pages; //amount of memory to be allocated for managing 4k page size DTLB

//TODO: followig not used for now. Meant to keep all color enums.
__thread enum color *color_list_small_4k;
__thread enum color *color_list_large_4k;

//array of pointers to free lists based on color
__thread  struct mblock **large_free_list_ptrs_4k;
__thread struct mblock **small_free_list_ptrs_4k;

/* Function prototypes */
static int init();
static void * allocate(unsigned int type, unsigned int nbytes, enum color color_needed);
static void * tlb_malloc_small(unsigned int nbytes, enum color color_needed);
static void tlb_free_small(void *addr, enum color color_needed);
static void * tlb_malloc_large(unsigned int nbytes, enum color color_needed);
static void tlb_free_large(void *addr, enum color color_needed);
static void * tlb_malloc_hugepage(unsigned int nbytes, enum color color_needed);
static enum color get_color(void *addr);
static int get_dtlb_set(void * addr);

void init_tlb_malloc_lib(unsigned int _4k_no_of_dtlb_sets, unsigned int _4k_dtlb_assoc, unsigned int _2m_dtlb_assoc, unsigned int _2m_no_of_dtlb_sets, unsigned int _4k_no_of_dtlb_sets_for_array, unsigned int _2m_no_of_dtlb_sets_for_array, unsigned int _assoc_multiplier)
{
	int i;

	//set all the thread specific global variables
	no_of_dtlb_sets_4k = _4k_no_of_dtlb_sets;
	no_of_dtlb_sets_for_array_4k = _4k_no_of_dtlb_sets_for_array;
	no_of_dtlb_sets_2m = _2m_no_of_dtlb_sets;
	no_of_dtlb_sets_for_array_2m =  _2m_no_of_dtlb_sets_for_array;
	dtlb_assoc_4k = _4k_dtlb_assoc;
	dtlb_assoc_2m = _2m_dtlb_assoc;
	assoc_multiplier = _assoc_multiplier;	

	//calculate values based on the parameters passed
	no_of_dtlb_sets_for_non_array_4k = no_of_dtlb_sets_4k - no_of_dtlb_sets_for_array_4k;
	memory_req_4k_pages = no_of_dtlb_sets_4k * dtlb_assoc_4k * ASSOC_MULTIPLIER_4K * PAGE_SIZE_4K;
	//memory_req_4k_pages = no_of_dtlb_sets_4k * dtlb_assoc_4k * assoc_multiplier * PAGE_SIZE_4K;

	//allocate space for free list and color list.
	large_free_list_ptrs_4k = (struct mblock **) malloc( no_of_dtlb_sets_for_array_4k * sizeof(struct mblock *));
	small_free_list_ptrs_4k = (struct mblock **) malloc( no_of_dtlb_sets_for_non_array_4k * sizeof(struct mblock *));

	//example if 4 way assoc, then space in tlb for array will have 4 colors, while non arrays will have colors equal to the number of sets dedicated to the non array part
	color_list_large_4k = (enum color *) malloc(dtlb_assoc_4k * sizeof(enum color));
	color_list_small_4k = (enum color *) malloc(no_of_dtlb_sets_for_non_array_4k * sizeof(enum color));

	//initialize free lists and color lists 
	for(i = 0; i < no_of_dtlb_sets_for_array_4k; i++)
	{
		large_free_list_ptrs_4k[i] = (struct mblock *) malloc( sizeof( struct mblock));
		large_free_list_ptrs_4k[i] -> next = NULL;
	}
	
	for(i = 0; i < no_of_dtlb_sets_for_non_array_4k; i ++)
	{
		small_free_list_ptrs_4k[i] = (struct mblock *) malloc( sizeof( struct mblock));
		small_free_list_ptrs_4k[i] -> next = NULL;
	}
	for(i = 0; i < dtlb_assoc_4k; i ++)
	{
		color_list_large_4k[i] = i;
	}
	for(i = 0; i < no_of_dtlb_sets_for_non_array_4k; i ++)
	{
		color_list_small_4k[i] = i;
	}
	init();
	printf("----------Printing large allocation lists------------\n");
	/*for( i = 0; i < dtlb_assoc_4k; i ++)
	{
		printf("COLOR %d\n",i);
		print(large_free_list_ptrs_4k, i);
		printf("\n\n");
	}*/
	printf("\n\n");
	printf("----------Printing small allocation lists-------------\n");
	for( i = 0; i < no_of_dtlb_sets_for_non_array_4k; i ++)
	{
		printf("COLOR %d\n",i);
		print(small_free_list_ptrs_4k, i);
		printf("\n\n");
	}
}

/*

This function returns a colored page in the virtual address space. The function allocates space in multiple of 16 bytes since the mblock struct needs 16 Bytes.
It also uses 4 bytes ( prceding the addresses returned to the user) to store the length of the allocation. This malloc works in a TLB friendly way and gives isoaltion for real time tasks. 

Parameters:
	nbytes - Specifies the number of bytes needed.
	color_needed - An enum which specifies the virtual page color that is needed.
Returns:
	A void pointer to the allocated space and NULL otherwise.

Notes:

	1)The function allocates space in multiple of 16 bytes since the mblock struct needs 16 Bytes.
	2)It also uses 4 bytes ( prceding the addresses returned to the user) to store the length of the allocation. 
	3) Small allocation means less than 4k allocation - Max sized allocation is 4092 Bytes - due to 4 Byte overhead for length
	4) Large allocation means more than 4k (4092B to be precise) allocation but less than 2M
	5) Huge allocation means greater than equal to 2M allocation

*/
void * tlb_malloc(unsigned int nbytes, enum color color_needed)
{
	static __thread int initialized = 0;
	static __thread void * mem_ptr = NULL;
	//TODO: Dont think we need this check since call to tlb_init is mandatory before using the tlb_malloc
	/*if(initialized == 0)
	{
		init();
		initialized = 1;
	}*/
	if(nbytes <= PAGE_SIZE_4K - BYTES_FOR_ALLOCATION_LENGTH)
	{
		mem_ptr = tlb_malloc_small(nbytes, color_needed);
	}
	else if( (nbytes > PAGE_SIZE_4K - BYTES_FOR_ALLOCATION_LENGTH) && (nbytes <= PAGE_SIZE_2M) )
	{
		mem_ptr = tlb_malloc_large(nbytes, color_needed);
	}
	else
	{
		mem_ptr = tlb_malloc_hugepage(nbytes, color_needed);
	}
	return mem_ptr;
}

/*

This function frees the memory allocated by tlb_malloc. 

Parameters:
	addr - the base address allocated by tlb_malloc
	color_needed - The color associated with addr.
Returns:
	void

Notes:
	1) The function calls tlb_free_large / tlb_free_small / tlb_free_huge based on the length of the allocated block  

*/
void tlb_free(void * addr, enum color color_needed)
{
	int success = 0;
	
	if( addr != NULL)
	{
		unsigned int * length_ptr = (unsigned int *) addr;
		//length_ptr --;
		length_ptr = length_ptr - 2;
		unsigned int length = *length_ptr;

		//int dtlb_set = get_dtlb_set(addr);
		//check if the set number is within the large section or the small section
		if( length > PAGE_SIZE_4K && length < PAGE_SIZE_2M)
		{
			tlb_free_large(addr, color_needed);
		}	
		else if( length > 0 && length <= PAGE_SIZE_4K) //if 4092 bytes are needed we actually allocate 4094 coz of 4 bytes overhead for keeping length
		{
			tlb_free_small(addr, color_needed);
		}
		else if (length >= PAGE_SIZE_2M)
		{
			//TODO: Add function call for tlb_free_huge
		}
	}	

}

/*

This function adds free memory to either the small / large / huge page free lists.

Parameters:
	type - 0 for small , 1 for large , 2 for huge list
	addr - the base address returned to the user by tlb_malloc
	color_needed - the color associated with addr

Returns:
	void

Notes: This function performs compaction.

*/
//type 0 for small, 1 for large 
static void free_memory(int type, void *addr, enum color color_needed)
{
	struct mblock * prev_ptr = NULL, *current_ptr = NULL;
	unsigned int * length_ptr = NULL;

	//the length of the allocation is kept in the preceding 4 Bytes
	length_ptr = (unsigned int *) addr;
	//length_ptr--;
	length_ptr = length_ptr - 2;
	unsigned int length = *length_ptr;
	
	void *original_addr_ptr = (void *) length_ptr; // this is the address we need to add back to our list
	if (type == 0)
	{
        	prev_ptr = small_free_list_ptrs_4k[color_needed];
        	current_ptr = small_free_list_ptrs_4k[color_needed] -> next;
		
	}
	else if(type == 1)
	{
        	prev_ptr = large_free_list_ptrs_4k[color_needed];
        	current_ptr = large_free_list_ptrs_4k[color_needed] -> next;
	}
	else if( type ==2 )
	{
		//TODO: add function call for huge page
	}

	//find the right place for the block to be added into the list. The lists are kept in increasing order of memory addresses. Add the end of traversal, prev_ptr will be pointing to memory
	// block whose address is less than original_addr_ptr while current_ptr will be pointing to the first free memory block whose address is greater than original_addr_ptr
	while( current_ptr != NULL)
	{
		if (  (void *) current_ptr <  original_addr_ptr)
		{
			prev_ptr = current_ptr;
			current_ptr = current_ptr -> next;
		}
		else
		{
			break;
		}
	}

	//check if compaction is possible and insert the block into the list 
	
	//check is the previous block and the block being freed can be coalesced
	if( ( (void *) (prev_ptr) + prev_ptr -> length) == original_addr_ptr && ( (original_addr_ptr + length) != (void *) current_ptr ))
	{
		prev_ptr -> length = prev_ptr -> length + length;
	}
	//check is the block being freed and the current block can be coalesced
	else if ( ( (original_addr_ptr + length) == (void *) current_ptr) && ( ( ( (void *) (prev_ptr ) ) + prev_ptr -> length) !=  original_addr_ptr))
	{
		struct mblock * mem_returned = (struct mblock *) original_addr_ptr;
		mem_returned -> length = length + current_ptr -> length;
		mem_returned -> next = current_ptr -> next;
		prev_ptr -> next = mem_returned;
	}
	//check if the block being freed lies excatly between the previous block and the current block
	else if ( ( ( original_addr_ptr + length) == (void *) current_ptr) && ( ( ( (void *) prev_ptr) + prev_ptr -> length) ==  original_addr_ptr))
	{
		prev_ptr -> length = prev_ptr -> length + length + current_ptr -> length;
		prev_ptr -> next = current_ptr -> next;
	}
	// no compaction possible - add to the right place
	else
	{
		struct mblock * mem_returned = ( struct mblock *) original_addr_ptr;
		prev_ptr -> next = mem_returned;
		mem_returned -> next = current_ptr;
		mem_returned -> length = length;
	}
	
}

/*

The following two functions, tlb_free_small and tlb_free_large have not been exposed for now but there may be a need to.

*/
static void tlb_free_small(void *addr, enum color color_needed)
{
	free_memory(0, addr,color_needed);
}

static void tlb_free_large(void *addr, enum color color_needed)
{
	free_memory(1, addr,color_needed);
}

/*

This function initializes the data structures needed for allocating and deallocating memory using tlb_malloc and tlb_free.

Notes:
	1) Allocates a huge per thread memory space. Allocates more than no_of_dtlb_sets * dtlb_assoc since we need to align to page boundary and also align to dtlb set zero
	2) Initializes the small / large / huge free lists

*/
static int init()
{
	unsigned long base_addr = 0;
	unsigned long rem = 0;
	unsigned long add_factor = 0;


	//1. Add 4k since we want to be page boundary alignment.
	//2. Add dtlb_sets * 4k since we want to align the addresses to dtlb set 0. This factor is a worst case one.
	memory_req_4k_pages = memory_req_4k_pages + PAGE_SIZE_4K + ( no_of_dtlb_sets_4k * PAGE_SIZE_4K);
	memory = malloc(memory_req_4k_pages); //need to add 4k since we will align to page boundary
	//printf("Thread %d Initial address: %p\n", pthread_self(), memory);
	//align memory to page size boundary
	if( (unsigned long) (memory) % PAGE_SIZE_4K != 0)
	{
		base_addr = (unsigned long) memory;
		rem = (base_addr) % PAGE_SIZE_4K;
		add_factor = PAGE_SIZE_4K - rem;
		memory = (void *) base_addr + add_factor;
	}

	int dtlb_set = get_dtlb_set(memory);
	memory = (void *)  ( memory +   ( ( no_of_dtlb_sets_4k - dtlb_set) * 4096 )  );
	//printf("Thread %d Base address: %p\n", pthread_self(),memory); 
	int no_of_pages_non_array_4k = no_of_dtlb_sets_for_non_array_4k * dtlb_assoc_4k * ASSOC_MULTIPLIER_4K;
	int i =0;
	void * memory_ptr = memory;

	if ( no_of_dtlb_sets_for_array_4k != 0)
	{
		//initialize large_free_list - list for array allocation
		for( i = 0; i < dtlb_assoc_4k - 1; i ++) //one less than assoc because of the addition factor for memory_ptr will go out for bounds in the last iteration otherwise
		{
			struct mblock * block_to_insert = (struct mblock *) memory_ptr;
			block_to_insert -> length = no_of_dtlb_sets_for_array_4k * PAGE_SIZE_4K;
			block_to_insert -> next = NULL;
			insert(large_free_list_ptrs_4k, i, block_to_insert);

			//first addtion: get past the large block for ith color
			memory_ptr = memory_ptr + ( no_of_dtlb_sets_for_array_4k * PAGE_SIZE_4K);

			//second addition: get past the area for small (less than one page) allocations 
			memory_ptr = memory_ptr + ( no_of_dtlb_sets_for_non_array_4k * PAGE_SIZE_4K);
		}

		//insert the last color
		struct mblock * block_to_insert = (struct mblock *) memory_ptr;
		block_to_insert -> length = no_of_dtlb_sets_for_array_4k * PAGE_SIZE_4K;
		block_to_insert -> next = NULL;
		insert(large_free_list_ptrs_4k, i, block_to_insert);

		//re adjust the memory pointer to point to the right starting address for non array part. In the above loop the last iteration increment throws the memory_ptr off.
		memory_ptr = memory + (no_of_dtlb_sets_for_array_4k * PAGE_SIZE_4K);
	}
	//initialize the small_free_list - list for non array allocations
	int pages_added = 0;
	while( pages_added != ( no_of_pages_non_array_4k) )
	{
		for( i = 0; i < no_of_dtlb_sets_for_non_array_4k; i ++)
		{
			//enum color page_color = get_color(memory_ptr);
			enum color page_color = i;
			struct mblock * block_to_insert = (struct mblock *) memory_ptr;
			block_to_insert -> length = PAGE_SIZE_4K;
			block_to_insert -> next = NULL;
			insert(small_free_list_ptrs_4k, page_color, block_to_insert);
			memory_ptr = memory_ptr + (1 * PAGE_SIZE_4K);
			pages_added ++;
		}
		memory_ptr = memory_ptr + ( no_of_dtlb_sets_for_array_4k * PAGE_SIZE_4K);	
	}
	struct color_ctrl_page *color_ctrl;

        //int i;
        //COLOR MALLOC PART
        color_ctrl = get_color_ctrl();
        if (!color_ctrl) {
                fprintf(stderr, "Could not map color control interface.\n");
                exit(EXIT_FAILURE);
        }
	int j;
	for(i=0;i<5 ;i++)
	{
		
		setup_pages(color_ctrl,i,1,0);
		struct mblock *ptr = small_free_list_ptrs_4k[i]->next;	
		struct mblock *nextptr = ptr->next;
		for(j=0;j< dtlb_assoc_4k;j++)
		{
			ptr=color_malloc(4096, (void*)ptr);   	
			ptr->length = PAGE_SIZE_4K;
			ptr->next = nextptr;
			ptr=ptr->next;
			if(ptr!=NULL)
			nextptr = ptr->next;
		}
	}
}

static enum color get_color(void * addr)
{
	//TODO: change this mask to be configured dynamically
	long addr_mask = 0x00000000F000;
	int color =  ( ( (long) addr & addr_mask) >> 12) % no_of_dtlb_sets_for_non_array_4k;
	return color_list_small_4k[color];
}

static int get_dtlb_set(void * addr)
{
	long addr_mask = 0x00000002F000;
	int color =  ( ( (long) addr & addr_mask) >> 12);
	return color;
}

/*

   This function allocates the memory required from either small / large / huge free list

Parameters:
type - 0 for small, 1 for large, 2 for huge list
nbytes - the number of bytes needed
color_needed - the color of the page needed
Returns:
a void pointer if allocation was successful and NULL otherwise.

 */
static void * allocate(unsigned int type, unsigned int nbytes, enum color color_needed)
{
	static __thread struct mblock *current_ptr, *prev_ptr, *leftover;
	if( type == 0 )
	{
		prev_ptr = small_free_list_ptrs_4k[color_needed];
		current_ptr = small_free_list_ptrs_4k[color_needed] -> next;
	}
	else if ( type == 1 )
	{
		prev_ptr = large_free_list_ptrs_4k[color_needed];
		current_ptr = large_free_list_ptrs_4k[color_needed] -> next;
	}
	else if( type == 2)
	{
		//TODO: add logic for huge free list
	}
	if(nbytes == 0 || current_ptr  == NULL)
		return NULL;

	nbytes = nbytes + sizeof(unsigned int);//to keep length of the block returned to user - this length info will be used while freeing
	nbytes = (unsigned int ) roundmb(nbytes);


	//using first-fit-policy
	while( current_ptr != NULL)
	{
		//if the bytes requested is exactly equal to a block
		if(current_ptr -> length == nbytes)
		{

			prev_ptr -> next = current_ptr -> next;
			unsigned int *block = (unsigned int *) current_ptr;
			*block = nbytes;
			//block ++;
			block = block + 2;
			current_ptr = (void *) block;
			return (void *) current_ptr;

		}
		//break the block into smaller one
		else if( current_ptr -> length > nbytes)
		{
			leftover =  (struct mblock *) ( (unsigned long) current_ptr + nbytes);
			prev_ptr -> next = leftover;
			leftover -> next = current_ptr -> next;
			leftover -> length = current_ptr -> length - nbytes;
			unsigned int *block = (unsigned int *) current_ptr;
			*block = nbytes;
			//block ++;
			block = block + 2;
			current_ptr = (void *) block;
			return (void *) current_ptr;
		}
		else
		{
			prev_ptr = current_ptr;
			current_ptr = current_ptr -> next;
		}
	}
	//if there are no free blocks which can satisfy the request
	return NULL;

}

/*

   The following function, tlb_malloc_small, tlb_malloc_large , tlb_malloc_huge are not exposed for now. But there may be a need to.

 */
static void * tlb_malloc_small(unsigned int nbytes, enum color color_needed)
{

	return allocate(0, nbytes, color_needed);	

}
static void * tlb_malloc_large(unsigned int nbytes, enum color color_needed)
{

	return allocate(1, nbytes, color_needed);	
}
static void * tlb_malloc_hugepage(unsigned int nbytes, enum color color_needed)
{
}
