#include "wp.h"
#include "xparameters.h"
#include "xio.h"
#include <string.h>
#include <stdio.h>
#include <stdlib.h>

#define BLOCK_SIZE 512
#define BIT_COUNTER_SIZE 256
#define BYTE_LIMIT 2048

void wp( Xuint32 baseaddr, unsigned char *buffer, int message_length, Xuint32 *output_hash);
void single_block_zeros(Xuint32 baseaddr);
void single_block_ones(Xuint32 baseaddr);
void byte_limit_blocks(Xuint32 baseaddr, unsigned char *message);

int main(void){
	Xuint32 baseaddr;
	int i;
	
	unsigned char message[] = "Przemyslaw Zalewski";

	//XCache_EnableICache(0xc0000001);
	//XCache_EnableDCache(0xc0000001);
	
	baseaddr = (Xuint32) XPAR_WP_CONFIG_1_0_BASEADDR;

	printf("-- Entering wp_test main()-\n");

	single_block_zeros(baseaddr);
	single_block_ones(baseaddr);
	byte_limit_blocks(baseaddr,message);
	
	printf("\n-- Exiting wp_test main()-\n");

	XCache_DisableDCache();
	XCache_DisableICache();	

	return XST_SUCCESS;
}

void wp( Xuint32 baseaddr, unsigned char *buffer, int message_length, Xuint32 *output_hash){

	int i,padded_message_bytes_length,num_blocks,num_ints;
	unsigned long long num_bits;
	Xuint32 *converted;
	int	Index;
	Xuint64 start,reset,data_num_blocks,wp_done,hash;
	unsigned char *padded_message;

	num_bits  = 8*message_length;
	num_blocks = num_bits / BLOCK_SIZE;
	if(num_bits % BLOCK_SIZE < BLOCK_SIZE - BIT_COUNTER_SIZE - 1){
		num_blocks += 1;
	}else{
		num_blocks += 2;
	}
	padded_message_bytes_length = num_blocks*(BLOCK_SIZE/8);

	padded_message = (unsigned char *)calloc(padded_message_bytes_length,sizeof(unsigned char));
	if(padded_message == NULL){
		printf("Allocation failed\n");
	}
	// prepend the actual message
	memcpy(padded_message,buffer,message_length);
	
	// append the '1' bit
	padded_message[message_length] = 0x80;

	// append the number of bits (as a 64 bit unsigned integer)
	for( i = 0; i < 8; i++){
		padded_message[padded_message_bytes_length-(i+1)] = (unsigned char)((num_bits >> i*8)   & 0xFF);
	}	

	num_ints = (BLOCK_SIZE/32)*num_blocks;
	converted 	= (Xuint32 *)malloc(num_ints);
	if(converted == NULL){
		printf("Allocation failed\n");
	}

	for(i=0; i<num_ints; i++)
	{
		converted[i]  = (int)padded_message[i*4+0] << 24;
		converted[i] |= (int)padded_message[i*4+1] << 16;
		converted[i] |= (int)padded_message[i*4+2] << 8;
		converted[i] |= padded_message[i*4+3];
	}

	/* Start setup fifo stuff */

	WP_mResetWriteFIFO(baseaddr);

	for ( Index = 0; Index < num_ints-1; Index+=2 )
	{
		WP_mWriteToFIFO(baseaddr, 0, converted[Index]);
		WP_mWriteToFIFO(baseaddr, 4, converted[Index+1]);
		//printf("Wrote %0.8X %0.8X\n",converted[Index],converted[Index+1]);
	}
	/* End setup fifo stuff */

	/* Setup WP stuff */
	reset.Upper = 0x00000000;   reset.Lower = 0x00000000;
	WP_mWriteSlaveReg3(baseaddr, 0, reset.Upper);
	WP_mWriteSlaveReg3(baseaddr, 4, reset.Lower);

	reset.Upper = 0x00000000;   reset.Lower = 0x00000001;
	WP_mWriteSlaveReg3(baseaddr, 0, reset.Upper);
	WP_mWriteSlaveReg3(baseaddr, 4, reset.Lower);

	start.Upper = 0x00000000;   start.Lower = 0x00000000;
	WP_mWriteSlaveReg0(baseaddr, 0, start.Upper);
	WP_mWriteSlaveReg0(baseaddr, 4, start.Lower);

	data_num_blocks.Upper = 0x00000000;   data_num_blocks.Lower = num_blocks;
	WP_mWriteSlaveReg1(baseaddr, 0, data_num_blocks.Upper);
	WP_mWriteSlaveReg1(baseaddr, 4, data_num_blocks.Lower);

	start.Upper = 0x00000000;   start.Lower = 0x00000001;
	WP_mWriteSlaveReg0(baseaddr, 0, start.Upper);
	WP_mWriteSlaveReg0(baseaddr, 4, start.Lower);

	wp_done.Upper = 0; wp_done.Lower = 0;

	while (wp_done.Lower != 0x00000001 ){
		//printf("Waiting for done \n");
		wp_done.Lower =  WP_mReadSlaveReg2(baseaddr, 4);
	}

	start.Upper = 0x00000000;   start.Lower = 0x00000000;
	WP_mWriteSlaveReg0(baseaddr, 0, start.Upper);
	WP_mWriteSlaveReg0(baseaddr, 4, start.Lower);

	/* Grab the final hash */
	for(i = 0; i < 16; i++){
		output_hash[i] = WP_mReadReg(baseaddr+WP_SLV_REG4_OFFSET,4*i);
	}
	free(padded_message);
	free(converted);	
}

void single_block_zeros(Xuint32 baseaddr){
	int i,j;
	unsigned char *buffer;
	Xuint32 digest[512/32];

	buffer = malloc(BLOCK_SIZE/8);
	if(buffer == NULL){
		printf("Allocation failed\n");
	}
	for( i = 0; i < BLOCK_SIZE/8; i++){
		buffer[i] = 0x00;
	}

	for( i = 1; i <= (int)(BLOCK_SIZE - BIT_COUNTER_SIZE - 1)/8; i++){
		wp(baseaddr,buffer, i, digest);
		for(j = 0; j < 16; j++){
			printf("%08X",digest[j]);
		}
	}
	free(buffer);
	
}
void single_block_ones(Xuint32 baseaddr){
	int i,j;
	unsigned char *buffer;
	Xuint32 digest[512/32];

	buffer = malloc(BLOCK_SIZE/8);
	if(buffer == NULL){
		printf("Allocation failed\n");
	}
	for( i = 0; i < BLOCK_SIZE/8; i++){
		buffer[i] = 0xFF;
	}

	for( i = 1; i <= (int)(BLOCK_SIZE - BIT_COUNTER_SIZE - 1)/8; i++){
		wp(baseaddr,buffer, i, digest);
		for(j = 0; j < 16; j++){
			printf("%08X",digest[j]);
		}
	}
	free(buffer);
	
}
void byte_limit_blocks(Xuint32 baseaddr, unsigned char *message){
	int i,j;
	unsigned char *buffer;
	Xuint32 digest[512/32];

	buffer = malloc(BYTE_LIMIT);
	if(buffer == NULL){
		printf("Allocation failed\n"); 
	}
	for( i = 0; i < BYTE_LIMIT; i++){
		buffer[i] = message[i % strlen(message)];
		//printf("%c",buffer[i]);
	}

	//for( i = 1; i <= BYTE_LIMIT; i++){	
	for( i = 1; i <= BYTE_LIMIT; i++){
	//	printf("i = %d\n",i);
		wp(baseaddr,buffer, i, digest);
		for(j = 0; j < 16; j++){
			printf("%08X",digest[j]);
		}
		printf("\n");	
	}
	free(buffer);
}