#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <sys/stat.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;
    char message[87];
    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;

    FILE *Handle=NULL;              /* File Handler */
    struct stat st = {0};            /* Dir Handler */

    NTRU_ENCRYPT_PARAM_SET_ID param_set_id;
    
    printf("------------------------------------------------\n");
    printf("Enter Message \t\t: ");
    // message = "Hello Bastard...";
    fgets(message, 86, stdin);

    printf("Your message is \t: %s\n",message);
    param_set_id = PARAM_SET_IDS[14]; /* 0 : 401; 1 : 449; 14 : 593;  */

    fprintf(stderr, "Testing parameter set \t: %s\n", ntru_encrypt_get_param_set_name(param_set_id));
    printf("------------------------------------------------\n");
    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 (stat("./keys", &st) == -1) {
        mkdir("./keys", 0700);
        printf("Key directory created...\n");

    }

    Handle=fopen("keys/key-593.pub", "wb");
    if(Handle!=NULL) {
        printf("-> Writing Pub Key...\n");
        fwrite(public_key, public_key_len, 1, Handle);
        printf("-> Pub Key written...\n");
        fclose(Handle);
    }

    Handle=fopen("keys/key-593.priv", "wb");
    if(Handle!=NULL) {
        printf("-> Writing Priv Key...\n");
        fwrite(private_key, private_key_len, 1, Handle);
        printf("-> Private Key written...\n");
        fclose(Handle);
    }

    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("------------------------------------------------\n");
    printf("Max message block \t: %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);


    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("Decryption result \t: %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\n",
            public_key_len, private_key_len, ciphertext_len);
    printf("------------------------------------------------\n");
    fprintf(stderr, "\n");
}