#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <string.h>
#include <time.h>
#include <math.h>

#include "rng.h"
#include "api.h"
#include "gmp.h"
#include "kaz_api.h"
#include "sha256.h"

void HashMsg(const unsigned char *msg, unsigned long long mlen, unsigned char buf[32])
{
    sha256_t hash;
	sha256_init(&hash);
	sha256_update(&hash, msg, mlen);
	sha256_final(&hash, buf);
}

void KAZ_DS_CRT(int size, mpz_t *x, mpz_t *modulus, mpz_t crt)
{
    mpz_t *c=malloc(size*sizeof(mpz_t));
    mpz_t u, prod;

    mpz_inits(u, prod, NULL);
    for(int i=0; i<size; i++) mpz_init(c[i]);

    mpz_set_ui(c[0], 0);

    for(int i=1; i<size; i++){
        mpz_set_ui(c[i], 1);

        for(int j=0; j<=i-1; j++){
            mpz_invert(u, modulus[j], modulus[i]);
            mpz_mul(c[i], c[i], u);
            mpz_mod(c[i], c[i], modulus[i]);
        }
    }

    mpz_set(u, x[0]);
    mpz_set(crt, u);

    for(int i=1; i<size; i++){
        mpz_sub(u, x[i], crt);
        mpz_mul(u, u, c[i]);
        mpz_mod(u, u, modulus[i]);
        mpz_set_ui(prod, 1);

        for(int j=0; j<=i-1; j++) mpz_mul(prod, prod, modulus[j]);

        mpz_mul(u, u, prod);
        mpz_add(crt, crt, u);
    }

    for(int i=0; i<size; i++) mpz_clear(c[i]);
    mpz_clears(u, prod, NULL);
}

int KAZ_DS_GET_PFactors(mpz_t input)
{
    mpz_t inp, initial;

    mpz_inits(inp, initial, NULL);

    mpz_set(inp, input);
    mpz_set_ui(initial, 2);

    int i=0;
    bool FLAGS=true;

    while(FLAGS==true){
        int count=0;

        while(mpz_divisible_p(inp, initial)>0){
            count++;
            mpz_divexact(inp, inp, initial);
        }

        int I=mpz_probab_prime_p(inp, 40);

        if(I>0){
            i++;
            FLAGS=false;
        }else{
            i++;
            mpz_nextprime(initial, initial);
        }
    }

    mpz_clears(inp, initial, NULL);

    return i;
}

void KAZ_DS_PFactors(mpz_t input, mpz_t *pfacs, mpz_t *qlist, int *elist)
{
    mpz_t inp, initial;

    mpz_inits(inp, initial, NULL);

    mpz_set(inp, input);
    mpz_set_ui(initial, 2);

    int i=0;
    bool FLAGS=true;

    while(FLAGS==true){
        int count=0;

        while(mpz_divisible_p(inp, initial)>0){
            count++;
            mpz_divexact(inp, inp, initial);
        }

        int I=mpz_probab_prime_p(inp, 40);

        if(I>0){
            mpz_set(pfacs[i], inp);
            mpz_set(qlist[i], inp);
            elist[i]=1;

            FLAGS=false;
        }else{
            mpz_pow_ui(pfacs[i], initial, count);
            mpz_set(qlist[i], initial);
            elist[i]=count;

            i++;
            mpz_nextprime(initial, initial);
        }
    }

    mpz_clears(inp, initial, NULL);
}

void KAZ_DS_RANDOM(int lb, int ub, mpz_t out)
{
	mpz_t lbound, ubound;

	gmp_randstate_t gmpRandState;
    gmp_randinit_mt(gmpRandState);
	mpz_inits(lbound, ubound, NULL);

	mpz_ui_pow_ui(lbound, 2, lb);
	mpz_ui_pow_ui(ubound, 2, ub);

	unsigned int sd=100000;

	do{
		gmp_randseed_ui(gmpRandState, rand()+sd);
		mpz_urandomb(out, gmpRandState, ub);
		sd+=1;
	}while((mpz_cmp(out, lbound) == -1) || (mpz_cmp(out, ubound) == 1));
}

void KAZ_DS_FILTER(mpz_t VQ, mpz_t V2, mpz_t GRg, mpz_t Q, mpz_t q, mpz_t SF2)
{
    mpz_t modulus, soln, g1, g2, g3, GRgqQ, qQ, tmp;

    mpz_inits(modulus, soln, g1, g2, g3, GRgqQ, qQ, tmp, NULL);

	mpz_mul(GRgqQ, GRg, q);
	mpz_mul(GRgqQ, GRgqQ, Q);
	mpz_mul(qQ, q, Q);

    int nGRgqQ=KAZ_DS_GET_PFactors(GRgqQ);

    mpz_t *pFactors=malloc(nGRgqQ*sizeof(mpz_t));
    mpz_t *p=malloc(nGRgqQ*sizeof(mpz_t));
    int *e=malloc(nGRgqQ*sizeof(int));

    for(int i=0; i<nGRgqQ; i++) mpz_init(pFactors[i]);
    for(int i=0; i<nGRgqQ; i++) mpz_init(p[i]);
    for(int i=0; i<nGRgqQ; i++) e[i]=0;

    KAZ_DS_PFactors(GRgqQ, pFactors, p, e);

    mpz_t *x=malloc(2*sizeof(mpz_t));
    mpz_t *y=malloc(2*sizeof(mpz_t));

    for(int i=0; i<2; i++) mpz_init(x[i]);
    for(int i=0; i<2; i++) mpz_init(y[i]);

    mpz_set_ui(SF2, 0);
    mpz_set_ui(modulus, 1);
	mpz_set_ui(soln, 0);

    for(int i=0; i<nGRgqQ; i++){
        mpz_gcd(g1, Q, pFactors[i]);
		mpz_gcd(g2, GRg, pFactors[i]);
		mpz_gcd(g3, qQ, pFactors[i]);
		mpz_mod(tmp, Q, g3);

		if(mpz_cmp_ui(g1, 1)!=0)
			mpz_set_ui(soln, 1);
		else if((mpz_cmp_ui(g1, 1)==0 && mpz_cmp_ui(g2, 1)!=0) && (mpz_cmp_ui(g3, 1)==0 || mpz_cmp_ui(tmp, 0)==0))
			mpz_mod(soln, VQ, g2);
		else if(mpz_cmp_ui(g1, 1)==0 && mpz_cmp_ui(g2, 1)==0)
			mpz_divexact(soln, V2, Q);

        mpz_set(x[0], SF2);
        mpz_set(x[1], soln);

        mpz_set(y[0], modulus);
        mpz_set(y[1], pFactors[i]);

        KAZ_DS_CRT(2, x, y, SF2);
        mpz_mul(modulus, modulus, pFactors[i]);
    }

    mpz_clears(modulus, soln, g1, g2, g3, GRgqQ, qQ, tmp, NULL);
}

void KAZ_DS_KeyGen(unsigned char *kaz_ds_verify_key, unsigned char *kaz_ds_sign_key)
{
    mpz_t GRg, phiGRg, phiphiGRg, q, Q, phiQ, qQ;
    mpz_t a, b, ALPHA, OMEGA1, V1, V2, SK, tmp1, tmp2;

    mpz_inits(GRg, phiGRg, phiphiGRg, q, Q, phiQ, qQ, NULL);
    mpz_inits(a, b, ALPHA, OMEGA1, V1, V2, SK, tmp1, tmp2, NULL);

    //1) Get all system parameters
    mpz_set_str(GRg, KAZ_DS_SP_GRg, 10);
    mpz_set_str(phiGRg, KAZ_DS_SP_PHIGRg, 10);
	mpz_set_str(phiphiGRg, KAZ_DS_SP_PHIPHIGRg, 10);
    mpz_set_str(q, KAZ_DS_SP_q, 10);
    mpz_set_str(Q, KAZ_DS_SP_Q, 10);
	mpz_set_str(phiQ, KAZ_DS_SP_PHIQ, 10);
	mpz_set_str(qQ, KAZ_DS_SP_qQ, 10);

	int k=KAZ_DS_SP_K;
	int LGg=KAZ_DS_SP_LGg;
	int LGRg=KAZ_DS_SP_LGRg;

    //1) Generate random a, ALPHA, OMEGA1
	KAZ_DS_RANDOM(LGg-2, LGg-1, a);
	mpz_nextprime(a, a);

	KAZ_DS_RANDOM(k+LGRg-1, k+LGRg, ALPHA);
	mpz_nextprime(ALPHA, ALPHA);

	KAZ_DS_RANDOM(LGg-2, LGg-1, OMEGA1);
	//mpz_nextprime(OMEGA1, OMEGA1);

	//2) Compute V1
	mpz_mod(V1, ALPHA, GRg);

	//3) Compute b
	mpz_mul(tmp1, OMEGA1, phiGRg);
	mpz_powm(b, a, phiphiGRg, tmp1);

	//4) Compute V2
	mpz_mul(tmp1, phiQ, b);
	mpz_powm(V2, ALPHA, tmp1, qQ);
	mpz_mul(V2, V2, Q);
	mpz_mod(V2, V2, qQ);

	//5) Compute Secret Signing SK
	mpz_mul(tmp1, phiQ, b);
	mpz_mul(tmp2, GRg, q);
	mpz_mul(tmp2, tmp2, Q);
	mpz_powm(SK, ALPHA, tmp1, tmp2);

    //6) Set kaz_ds_sign_key=(SK) & kaz_ds_verify_key=(V1, V2)
    size_t SKSIZE=mpz_sizeinbase(SK, 16);
	size_t V1SIZE=mpz_sizeinbase(V1, 16);
	size_t V2SIZE=mpz_sizeinbase(V2, 16);

	unsigned char *SKBYTE=(unsigned char*) malloc(SKSIZE*sizeof(unsigned char));
	mpz_export(SKBYTE, &SKSIZE, 1, sizeof(char), 0, 0, SK);

	unsigned char *V1BYTE=(unsigned char*) malloc(V1SIZE*sizeof(unsigned char));
	mpz_export(V1BYTE, &V1SIZE, 1, sizeof(char), 0, 0, V1);

	unsigned char *V2BYTE=(unsigned char*) malloc(V2SIZE*sizeof(unsigned char));
	mpz_export(V2BYTE, &V2SIZE, 1, sizeof(char), 0, 0, V2);

	for(int i=0; i<CRYPTO_SECRETKEYBYTES; i++) kaz_ds_sign_key[i]=0;

	int je=CRYPTO_SECRETKEYBYTES-1;
	for(int i=SKSIZE-1; i>=0; i--){
		kaz_ds_sign_key[je]=SKBYTE[i];
		je--;
	}

	for(int i=0; i<CRYPTO_PUBLICKEYBYTES; i++) kaz_ds_verify_key[i]=0;

	je=CRYPTO_PUBLICKEYBYTES-1;
	for(int i=V2SIZE-1; i>=0; i--){
		kaz_ds_verify_key[je]=V2BYTE[i];
		je--;
	}

	je=CRYPTO_PUBLICKEYBYTES-KAZ_DS_V2BYTES-1;
	for(int i=V1SIZE-1; i>=0; i--){
		kaz_ds_verify_key[je]=V1BYTE[i];
		je--;
	}

	mpz_clears(GRg, phiGRg, phiphiGRg, q, Q, phiQ, qQ, NULL);
    mpz_clears(a, b, ALPHA, OMEGA1, V1, V2, SK, tmp1, tmp2, NULL);
}

int KAZ_DS_SIGNATURE(unsigned char *signature, unsigned long long *signlen, const unsigned char *m, unsigned long long mlen, const unsigned char *sk)
{
    mpz_t GRg, phiGRg, phiphiGRg, q, Q, qQ, phiqQ, SK;
    mpz_t hashValue, r, OMEGA2, BETA, S, tmp1, tmp2;

    mpz_inits(GRg, phiGRg, phiphiGRg, q, Q, qQ, phiqQ, SK, NULL);
    mpz_inits(hashValue, r, OMEGA2, BETA, S, tmp1, tmp2, NULL);

    //1) Get all system parameters
	mpz_set_str(GRg, KAZ_DS_SP_GRg, 10);
	mpz_set_str(phiGRg, KAZ_DS_SP_PHIGRg, 10);
	mpz_set_str(phiphiGRg, KAZ_DS_SP_PHIPHIGRg, 10);
	mpz_set_str(q, KAZ_DS_SP_q, 10);
	mpz_set_str(Q, KAZ_DS_SP_Q, 10);
	mpz_set_str(qQ, KAZ_DS_SP_qQ, 10);
	mpz_set_str(phiqQ, KAZ_DS_SP_PHIqQ, 10);

	int LGg=KAZ_DS_SP_LGg;

    //2) Get kaz_ds_sign_key=(SK)
	unsigned char *SKBYTE=(unsigned char*) malloc((KAZ_DS_SKBYTES)*sizeof(unsigned char));

	for(int i=0; i<KAZ_DS_SKBYTES; i++) SKBYTE[i]=0;
	for(int i=0; i<KAZ_DS_SKBYTES; i++){SKBYTE[i]=sk[i];}

	mpz_import(SK, KAZ_DS_SKBYTES, 1, sizeof(char), 0, 0, SKBYTE);

	//3) Compute HASHValue(m)
	unsigned char buf[CRYPTO_BYTES]={0};
	HashMsg(m, mlen, buf);

	mpz_import(hashValue, CRYPTO_BYTES, 1, sizeof(char), 0, 0, buf);
	mpz_nextprime(hashValue, hashValue);

	//4) Generate random r & OMEGA2
    KAZ_DS_RANDOM(LGg-2, LGg-1, r);
	mpz_nextprime(r, r);

	KAZ_DS_RANDOM(LGg-2, LGg-1, OMEGA2);

	//5) Compute BETA
	mpz_mul(tmp1, OMEGA2, phiGRg);
	mpz_powm(BETA, r, phiphiGRg, tmp1);

    //6) Compute Signature
	mpz_mul(tmp1, phiqQ, BETA);
	mpz_mul(tmp2, GRg, q);
	mpz_mul(tmp2, tmp2, Q);//gmp_printf("SK=%Zd\n", tmp2);
	mpz_powm(S, hashValue, tmp1, tmp2);
	mpz_mul(S, S, SK);
	mpz_mod(S, S, tmp2);

    //6) Set signature=(S, m)
    size_t SSIZE=mpz_sizeinbase(S, 16);

	unsigned char *SBYTE=(unsigned char*) malloc(SSIZE*sizeof(unsigned char));
	mpz_export(SBYTE, &SSIZE, 1, sizeof(char), 0, 0, S);

	for(int i=0; i<mlen+KAZ_DS_SBYTES; i++) signature[i]=0;

	int je=mlen+KAZ_DS_SBYTES-1;
	for(int i=mlen-1; i>=0; i--){
		signature[je]=m[i];
		je--;
	}

	je=KAZ_DS_SBYTES-1;
	for(int i=SSIZE-1; i>=0; i--){
		signature[je]=SBYTE[i];
		je--;
	}

	*signlen=mlen+KAZ_DS_SBYTES;

    free(SBYTE);

	mpz_clears(GRg, phiGRg, phiphiGRg, q, Q, qQ, phiqQ, SK, NULL);
    mpz_clears(hashValue, r, OMEGA2, BETA, S, tmp1, tmp2, NULL);

    return 0;
}

int KAZ_DS_VERIFICATION(unsigned char *m, unsigned long long *mlen, const unsigned char *sm, unsigned long long smlen, const unsigned char *pk)
{
    mpz_t N, g, Gg, R, GRg, q, Q, phiQ, GRgQ, GRgqQ, qQ, phiqQ, hashValue, V1, V2, S;
    mpz_t tmp, tmp2, SF1, SF2, W0, W1, W2, W3, W4, VQ, y, y1, y2;

    mpz_inits(N, g, Gg, R, GRg, q, Q, phiQ, GRgQ, GRgqQ, qQ, phiqQ, hashValue, V1, V2, S, NULL);
    mpz_inits(tmp, tmp2, SF1, SF2, W0, W1, W2, W3, W4, VQ, y, y1, y2, NULL);

    //1) Get all system parameters
    mpz_set_str(N, KAZ_DS_SP_N, 10);
    mpz_set_str(g, KAZ_DS_SP_G, 10);
    mpz_set_str(Gg, KAZ_DS_SP_Gg, 10);
    mpz_set_str(R, KAZ_DS_SP_R, 10);
    mpz_set_str(GRg, KAZ_DS_SP_GRg, 10);
    mpz_set_str(q, KAZ_DS_SP_q, 10);
    mpz_set_str(Q, KAZ_DS_SP_Q, 10);
    mpz_set_str(phiQ, KAZ_DS_SP_PHIQ, 10);
	mpz_set_str(qQ, KAZ_DS_SP_qQ, 10);
	mpz_set_str(phiqQ, KAZ_DS_SP_PHIqQ, 10);

    //2) Get kaz_ds_verify_key=(V1, V2)
	unsigned char *V1BYTE=(unsigned char*) malloc((KAZ_DS_V1BYTES)*sizeof(unsigned char));
	unsigned char *V2BYTE=(unsigned char*) malloc((KAZ_DS_V2BYTES)*sizeof(unsigned char));

	for(int i=0; i<KAZ_DS_V1BYTES; i++) V1BYTE[i]=0;
	for(int i=0; i<KAZ_DS_V2BYTES; i++) V2BYTE[i]=0;

	for(int i=0; i<KAZ_DS_V1BYTES; i++){V1BYTE[i]=pk[i];}
	for(int i=0; i<KAZ_DS_V2BYTES; i++){V2BYTE[i]=pk[i+KAZ_DS_V1BYTES];}

	mpz_import(V1, KAZ_DS_V1BYTES, 1, sizeof(char), 0, 0, V1BYTE);
	mpz_import(V2, KAZ_DS_V2BYTES, 1, sizeof(char), 0, 0, V2BYTE);

    //3) Get signature=(S, m)
    int len=smlen-KAZ_DS_SBYTES;

    unsigned char *SBYTE=(unsigned char*) malloc(KAZ_DS_SBYTES*sizeof(unsigned char));
	unsigned char *MBYTE=(unsigned char*) malloc(len*sizeof(unsigned char));

	for(int i=0; i<KAZ_DS_SBYTES; i++) SBYTE[i]=0;
	for(int i=0; i<len; i++) MBYTE[i]=0;

	for(int i=0; i<KAZ_DS_SBYTES; i++){SBYTE[i]=sm[i];}
	for(int i=0; i<len; i++){MBYTE[i]=sm[i+KAZ_DS_SBYTES];}

    mpz_import(S, KAZ_DS_SBYTES, 1, sizeof(char), 0, 0, SBYTE);

	//4) Compute the hash value of the message
    unsigned char buf[CRYPTO_BYTES]={0};
    HashMsg(MBYTE, len, buf);

    mpz_import(hashValue, CRYPTO_BYTES, 1, sizeof(char), 0, 0, buf);
    mpz_nextprime(hashValue, hashValue);

    //5) Filtering Procedures
	mpz_mul(GRgQ, GRg, Q);
	mpz_powm(y, V1, phiQ, GRgQ);
	mpz_powm(tmp, hashValue, phiqQ, GRgQ);
	mpz_mul(y, y, tmp);
	mpz_mod(y, y, GRgQ);

	mpz_t *x=malloc(2*sizeof(mpz_t));
    mpz_t *modulus=malloc(2*sizeof(mpz_t));

    for(int i=0; i<2; i++) mpz_init(x[i]);
	for(int i=0; i<2; i++) mpz_init(modulus[i]);

	mpz_divexact(x[0], V2, Q);
	mpz_set(x[1], y);

	mpz_set(modulus[0], q);
	mpz_set(modulus[1], GRgQ);

	KAZ_DS_CRT(2, x, modulus, SF1);

    mpz_powm(VQ, V1, phiQ, GRg);
	mpz_powm(tmp, hashValue, phiqQ, GRg);
	mpz_mul(VQ, VQ, tmp);
	mpz_mod(VQ, VQ, GRg);

    KAZ_DS_FILTER(VQ, V2, GRg, Q, q, SF2);

    //FILTER 1
	mpz_mul(GRgqQ, GRg, q);
	mpz_mul(GRgqQ, GRgqQ, Q);

    mpz_mod(tmp, S, GRgqQ);
    mpz_sub(W0, tmp, S);

    if(mpz_cmp_ui(W0, 0)!=0){
        printf("Filter 1...\n");
        return -4;
    }

	//FILTER 2
    mpz_sub(W1, tmp, SF1);

    if(mpz_cmp_ui(W1, 0)==0){
        printf("Filter 2...\n");
        return -4;
    }

	//FILTER 3
    mpz_sub(W2, tmp, SF2);

    if(mpz_cmp_ui(W2, 0)==0){
        printf("Filter 3...\n");
        return -4;
    }

    //FILTER 4
    mpz_mul(W3, Q, S);
	mpz_mod(W3, W3, qQ);
	mpz_sub(W4, W3, V2);

    if(mpz_cmp_ui(W4, 0)!=0){
        printf("Filter 4...\n");
        return -4;
    }

    //6) Verifying Procedures
    mpz_powm(tmp, R, S, Gg);
    mpz_powm(y1, g, tmp, N);

    mpz_powm(tmp, V1, phiQ, GRg);
	mpz_powm(tmp2, hashValue, phiqQ, GRg);
    mpz_mul(tmp, tmp, tmp2);
    mpz_mod(tmp, tmp, GRg);
    mpz_powm(tmp, R, tmp, Gg);
    mpz_powm(y2, g, tmp, N);

    if(mpz_cmp(y1, y2)!=0)
        return -4;

    memcpy(m, MBYTE, len);
    *mlen=len;

    mpz_clears(N, g, Gg, R, GRg, q, Q, phiQ, GRgQ, GRgqQ, qQ, phiqQ, hashValue, V1, V2, S, NULL);
    mpz_clears(tmp, tmp2, SF1, SF2, W0, W1, W2, W3, W4, VQ, y, y1, y2, NULL);

	return 0;
}
