#include <stdlib.h>
#include <string.h>
#include <stdio.h>

#define DIGEST 32 // 256-bit digest in bytes.

#define ROUNDS 24 // Number of KECCAK rounds to perform for SHA3-256.
#define WIDTH 200 // 1600-bit width in bytes.
#define LANES 25 // The state is an unrolled 5x5 array of 64-bit lanes.
#define RATE 136 // 1600-bit width - 512-bit capacity in bytes.

#define ROTL64(x, y) (((x) << (y)) | ((x) >> (64 - (y))))

typedef unsigned long long uint64_t;
typedef unsigned char uint8_t;

union state {
	uint64_t words[LANES];
	uint8_t bytes[WIDTH];
};

static void theta(void *arg) {
	union state *state = (union state*) arg;
	uint64_t C[5] = { 0 };
	uint64_t D[5] = { 0 };
	for (int i = 0; i < 5; i++) {
		C[i] = state->words[i];
		C[i] ^= state->words[i + 5];
		C[i] ^= state->words[i + 10];
		C[i] ^= state->words[i + 15];
		C[i] ^= state->words[i + 20];
	}
	for (int i = 0; i < 5; i++) {
		D[i] = C[(i + 4) % 5] ^ ROTL64(C[(i + 1) % 5], 1);
	}
	for (int i = 0; i < 5; i++) {
		for (int j = 0; j < 25; j += 5) {
			state->words[i + j] ^= D[i];
		}
	}
}

static void rho(void *arg) {
	union state *state = (union state*) arg;
	const int rotations[25] = {
		0, 1, 62, 28, 27,
		36, 44, 6, 55, 20,
		3, 10, 43, 25, 39,
		41, 45, 15, 21, 8,
		18, 2, 61, 56, 14
	};
	for (int i = 0; i < 25; i++) {
		state->words[i] = ROTL64(state->words[i], rotations[i]);
	}
}

static void pi(void *arg) {
	union state *state = (union state*) arg;
	const int switcheroo[25] = {
		0, 10, 20, 5, 15,
		16, 1, 11, 21, 6,
		7, 17, 2, 12, 22,
		23, 8, 18, 3, 13,
		14, 24, 9, 19, 4
	};
	uint64_t temp[25] = { 0 };
	for (int i = 0; i < 25; i++) {
		temp[i] = state->words[i];
	}
	for (int i = 0; i < 25; i++) {
		state->words[switcheroo[i]] = temp[i];
	}
}

static void chi(void *arg) {
	union state *state = (union state*) arg;
	uint64_t temp[5] = { 0 };
	for (int j = 0; j < 25; j += 5) {
		for (int i = 0; i < 5; i++) {
			temp[i] = state->words[i + j];
		}
		for (int i = 0; i < 5; i++) {
			state->words[i + j] ^= (~temp[(i + 1) % 5]) & temp[(i + 2) % 5];
		}
	}
}

static void iota(const uint8_t round, void *arg) {
	union state *state = (union state*) arg;
	const uint64_t constants[24] = {
		0x0000000000000001, 0x0000000000008082, 0x800000000000808a,
		0x8000000080008000, 0x000000000000808b, 0x0000000080000001,
		0x8000000080008081, 0x8000000000008009, 0x000000000000008a,
		0x0000000000000088, 0x0000000080008009, 0x000000008000000a,
		0x000000008000808b, 0x800000000000008b, 0x8000000000008089,
		0x8000000000008003, 0x8000000000008002, 0x8000000000000080,
		0x000000000000800a, 0x800000008000000a, 0x8000000080008081,
		0x8000000000008080, 0x0000000080000001, 0x8000000080008008
	};
	state->words[0] ^= constants[round];
}

static void keccak(void *arg) {
	union state *state = (union state*) arg;
	for (int round = 0; round < ROUNDS; round++) {
		theta(state); rho(state); pi(state); chi(state); iota(round,state);
	}
}

static int absorb(const uint64_t len, const uint8_t data[len], int absorbed, void *arg) {
	union state *state = (union state*) arg;
	for (uint64_t i = 0; i < len; i++) {
		state->bytes[absorbed++] ^= data[i];
		if (absorbed == RATE) {
			keccak(state);
			absorbed = 0;
		}
	}
	return absorbed;
}

static void squeeze(uint8_t digest[DIGEST], int padpoint, void *arg) {
	union state *state = (union state*) arg;
	state->bytes[padpoint] ^= 0x06;
	state->bytes[RATE - 1] ^= 0x80;
	keccak(state);
	for (int i = 0; i < DIGEST; i++) {
		digest[i] = state->bytes[i];
	}
}

void sha3_hash(uint8_t digest[DIGEST], const uint64_t len, const uint8_t data[len]) {
	union state *state = malloc(sizeof(union state));
	memset(state,'\0',sizeof(union state)); // absolutely necessary
	int absorbed = 0;
	int padpoint = absorb(len, data, absorbed, state);
	squeeze(digest, padpoint, state);
	free(state); state = NULL;
}

void printDig(uint8_t* cPtr, uint64_t len){
	uint64_t i=0;
	uint8_t *p=cPtr;
	for (i=0;i<len;i++){
		printf("%x",*p);
		p++;
	}
}


#define DATA_LEN 32

int main(){
	
	//Dummy digest
	uint8_t dig[DIGEST] = { 
		0x0,0x1,0x2,0x3,0x4,0x5,0x6,0x7,0x8,0x9,
		0x0,0x1,0x2,0x3,0x4,0x5,0x6,0x7,0x8,0x9,
		0x0,0x1,0x2,0x3,0x4,0x5,0x6,0x7,0x8,0x9,
		0xA,0xB
	};

	uint8_t data[DATA_LEN]={
		0x0,0x1,0x2,0x3,0x4,0x5,0x6,0x7,0x8,0x9,
		0x0,0x1,0x2,0x3,0x4,0x5,0x6,0x7,0x8,0x9,
		0x0,0x1,0x2,0x3,0x4,0x5,0x6,0x7,0x8,0x9,
		0xA,0xC
	};

	sha3_hash(dig,DATA_LEN,data);

	printDig(dig,DATA_LEN);

	return 0;
}