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

#include "ntru_crypto.h"
#include "ntru_crypto_drbg.h"
#include "test_common.h"

int main() 
{
	
    uint8_t *public_key;
    uint8_t *private_key;
    uint8_t *message;
    uint8_t *ciphertext;
    uint8_t *plaintext;

    uint16_t max_msg_len;
    uint16_t public_key_len;          /* no. of octets in public key */
    uint16_t private_key_len;         /* no. of octets in private key */
    uint16_t ciphertext_len;          /* no. of octets in ciphertext */
    uint16_t plaintext_len;           /* no. of octets in plaintext */
    DRBG_HANDLE drbg;                 /* handle for instantiated DRBG */
    uint32_t rc;                      /* return code */

    clock_t clk;

    NTRU_ENCRYPT_PARAM_SET_ID param_set_id;

    message = "AAAAAAAAAAAAAAAAAAAA";
    printf("Mess %s\n",message);
    param_set_id = PARAM_SET_IDS[0];

    fprintf(stderr, "Testing parameter set %s... \n", ntru_encrypt_get_param_set_name(param_set_id));
    fflush (stderr);
    rc = ntru_crypto_drbg_external_instantiate(
                                       (RANDOM_BYTES_FN) &randombytes, &drbg);

    if (rc != DRBG_OK)
    {
      fprintf(stderr,"\tError: An error occurred instantiating the DRBG\n");
    }

    rc = ntru_crypto_ntru_encrypt_keygen(drbg, param_set_id, &public_key_len,
                                         NULL, &private_key_len, NULL);
    if (rc != NTRU_OK)
    {
      ntru_crypto_drbg_uninstantiate(drbg);
      fprintf(stderr,"\tError: An error occurred getting the key lengths\n");
    }

    public_key = (uint8_t *)malloc(public_key_len * sizeof(uint8_t));
    private_key = (uint8_t *)malloc(private_key_len * sizeof(uint8_t));

    clk = clock();
    rc = ntru_crypto_ntru_encrypt_keygen(drbg, param_set_id, &public_key_len,
       					 public_key,
                                         &private_key_len,
                                         private_key);
    clk = clock() - clk;

    if (rc != NTRU_OK)
    {
      ntru_crypto_drbg_uninstantiate(drbg);
      free(public_key);
      free(private_key);
      fprintf(stderr,"\tError: An error occurred during key generation\n");
    }


    rc = ntru_crypto_ntru_encrypt(drbg, public_key_len, public_key, 0, NULL,
                                  &ciphertext_len, NULL);
    if (rc != NTRU_OK)
    {
      fprintf(stderr,"\tError: Bad public key");
    }

    rc = ntru_crypto_ntru_decrypt(private_key_len, private_key, 0, NULL,
                                  &max_msg_len, NULL);
    if (rc != NTRU_OK)
    {
      fprintf(stderr,"\tError: Bad private key");
    }

    //message = (uint8_t *) malloc(max_msg_len * sizeof(uint8_t));
    printf("Max Len : %d\n", max_msg_len * sizeof(uint8_t));

    ciphertext = (uint8_t *) malloc(ciphertext_len * sizeof(uint8_t));

    plaintext = (uint8_t *) malloc(max_msg_len * sizeof(uint8_t));

    plaintext_len = max_msg_len;

    //randombytes(message, max_msg_len);
    //randombytes(ciphertext, ciphertext_len);
    //randombytes(plaintext, plaintext_len);

    printf("Current Message: %s\n", message);

    clk = clock();
    rc = ntru_crypto_ntru_encrypt(drbg, public_key_len, public_key,
            max_msg_len, message, &ciphertext_len, ciphertext);
    clk = clock() - clk;
    if (rc != NTRU_OK){
      fprintf(stderr, "\tError: Encryption error %x\n", rc);
    }
    
    //printf("Cipher %s\n", ciphertext);


    clk = clock();
    rc = ntru_crypto_ntru_decrypt(private_key_len, private_key,
            ciphertext_len, ciphertext,
            &plaintext_len, plaintext);
    clk = clock() - clk;
    if (rc != NTRU_OK)
    {
      fprintf(stderr, "\tError: Decryption error %x\n", rc);
    }

    printf("Plain %s\n", plaintext);

    if(plaintext_len != max_msg_len || memcmp(plaintext,message,max_msg_len))
    {
      fprintf(stderr,
        "\tError: Decryption result does not match original plaintext\n");
    }

    ntru_crypto_drbg_uninstantiate(drbg);
    free(message);
    free(public_key);
    free(private_key);
    free(plaintext);
    free(ciphertext);

    fprintf(stderr, "pk %d, sk %d, ct %d bytes",
            public_key_len, private_key_len-public_key_len, ciphertext_len);
    fprintf(stderr, "\n");
}