#include <stdint.h>
#include "pke.h"
#include "fips202.h"
#include "parameter.h"
#include <string.h>
#include "rand.h"
#include "kem.h"
#include "poly.h"

unsigned long long  kem_get_pk_byts()
{
	return KEM_PUBLICKEYBYTES;
}
unsigned long long  kem_get_sk_byts()
{
	return KEM_SECRETKEYBYTES;
}
unsigned long long  kem_get_ss_byts()
{
	return KEM_KEYBYTES;

}
unsigned long long  kem_get_ct_byts()
{

	return KEM_CIP_BYTES;
}


int kem_keygen(unsigned char *pk, unsigned long long *pk_byts, unsigned char *sk, unsigned long long *sk_byts)
{
	int i;

	pke_keygen(pk,pk_byts,sk,sk_byts);

	sk += PKE_SECRETKEYBYTES;
	for (i = 0; i < PKE_PUBLICKEYBYTES; i++)
		sk[i] = pk[i];  //sk||pk

	sk += PKE_PUBLICKEYBYTES;
	shake256(sk, PKE_SEED_BYTES, pk, PKE_PUBLICKEYBYTES); //sk||pk||hash(pk)

	sk += PKE_SEED_BYTES;
	rand_byts(PKE_MESS_BYTES, sk);//sk||pk||hash(pk)||r //r=128

	*pk_byts = KEM_PUBLICKEYBYTES;
	*sk_byts = KEM_SECRETKEYBYTES;
	
	return 0;

}


int kem_enc(unsigned char *pk, unsigned long long pk_byts, unsigned char  *ss, unsigned long long* ss_byts,
	unsigned char *ct, unsigned long long *ct_byts)
{


	unsigned char coins[PKE_MESS_BYTES+ PKE_SEED_BYTES];
	unsigned char buf[PKE_SEED_BYTES + N / 8]; //buf=message||32×Ö½Ú

	int i;
	unsigned char mlen = PKE_MESS_BYTES;


	rand_byts(N / 8, buf);
	shake256(buf, N / 8, buf, N / 8);//message   


	shake256(buf + N / 8, PKE_SEED_BYTES, pk, PKE_PUBLICKEYBYTES); // u||hash(pk) 
	shake256(coins,PKE_MESS_BYTES+ PKE_SEED_BYTES, buf, PKE_SEED_BYTES + N / 8);//coins=hash(u||hash(pk))=K||s K:128

	pke_enc(pk, pk_byts, buf, mlen, ct, ct_byts, coins + PKE_MESS_BYTES);

	shake256(coins + PKE_MESS_BYTES, PKE_SEED_BYTES, ct, KEM_CIP_BYTES); //K||hash(c)

	shake256(ss, KEM_KEYBYTES, coins, PKE_MESS_BYTES + PKE_SEED_BYTES);  //ss=hash(K||hash(c))                          

	*ss_byts = KEM_KEYBYTES;
	*ct_byts = KEM_CIP_BYTES;

	return 0;





}


int kem_dec(unsigned char *sk, unsigned long long sk_byts,
	unsigned char *ct, unsigned long long ct_byts,
	unsigned char *ss, unsigned long long *ss_byts)
{



	int i, fail;
	unsigned char c_cmp[KEM_CIP_BYTES];
	unsigned char buf[PKE_SEED_BYTES + N / 8];
	unsigned char coins[PKE_MESS_BYTES + PKE_SEED_BYTES];                                                /* Will contain key, coins, qrom-hash */
	unsigned char *pk = sk + PKE_SECRETKEYBYTES;
	unsigned long long cpa_pklen = PKE_PUBLICKEYBYTES, cpa_sklen = PKE_SECRETKEYBYTES, cpa_clen = PKE_CIP_BYTES, mlen = PKE_MESS_BYTES;

	pke_dec(sk, cpa_sklen, ct, cpa_clen, buf, &mlen);

	for (i = 0; i < PKE_SEED_BYTES; i++)
		buf[N / 8 + i] = sk[KEM_SECRETKEYBYTES - PKE_SEED_BYTES- PKE_MESS_BYTES + i];//buf=m||hash(pk)

	shake256(coins, PKE_MESS_BYTES + PKE_SEED_BYTES, buf, N / 8 + PKE_SEED_BYTES);//K'||s'=hash(m||hash(pk))

	pke_enc(pk, cpa_pklen, buf, mlen, c_cmp, &cpa_clen, coins + PKE_MESS_BYTES); //c'


	fail = verify(ct, c_cmp, KEM_CIP_BYTES);//fail=0 success

	shake256(coins + PKE_MESS_BYTES, PKE_SEED_BYTES, ct, KEM_CIP_BYTES);//K'||hash(c)

	cmov(coins, sk + KEM_SECRETKEYBYTES - PKE_MESS_BYTES, PKE_MESS_BYTES, fail);//fail=0; ss=hash(K||hash(c)) success else  ss=hash(r||hash(c'))

	shake256(ss, KEM_KEYBYTES, coins, PKE_MESS_BYTES + PKE_SEED_BYTES);


	*ss_byts = KEM_KEYBYTES;
		
	return -fail;





}