#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdbool.h>
#include <assert.h>
#include "crypto_aead.h"
#include "api.h"

#ifndef KEY_LENGTH
#define KEY_LENGTH 16
#endif

#ifndef NSEC_LENGTH
#define NSEC_LENGTH 0
#endif

#ifndef NPUB_LENGTH
#define NPUB_LENGTH 0
#endif

#define FPRINTF_HEX(fd, m, mlen) for (unsigned long long int i = 0; i < mlen; i++) fprintf(fd, "%02x ", m[i]);

unsigned char *c = NULL;
unsigned long long clen = 0;
unsigned char *m = NULL;
unsigned long long mlen = 0;
unsigned char *ad = NULL;
unsigned long long adlen = 0;
unsigned char *nsec = NULL;

unsigned long long nslen = NSEC_LENGTH;
unsigned char *npub = NULL;
unsigned long long nplen = NPUB_LENGTH;
unsigned char *k = NULL;
unsigned long long klen = KEY_LENGTH;

static void recv_variable(unsigned char **target, unsigned long long *lenp) {
    if (*target != NULL) {
        free(*target);
    }
    
    uint32_t len;
    if (1 != fread(&len, sizeof(len), 1, stdin)) {
        fprintf(stderr, "ERROR: couldn't read length\r\n");
        exit(1);
    }
    *lenp = len;

    *target = malloc(len);

    if (NULL == *target) {
        fprintf(stderr, "ERROR: couldn't malloc %u bytes\r\n", len);
        exit(2);
    }

    if (*lenp == 0) {
        return;
    }

    if (1 != fread(*target, *lenp, 1, stdin)) {
        fprintf(stderr, "ERROR: didn't read %u bytes of data\r\n", len);
        exit(1);
    }
}


static void send_variable(unsigned char *target, unsigned long long len) {
    uint32_t olen = len;
    if (1 != fwrite(&olen, sizeof(olen), 1, stdout)) {
        fprintf(stderr, "ERROR: didn't write length\r\n");
        exit(1);
    }

    if (len == 0) {
        return;
    }

    if (1 != fwrite(target, olen, 1, stdout)) {
        fprintf(stderr, "ERROR: didn't write %llu bytes of data\r\n", len);
        exit(1);
    }
}


int main() {
    int res;
    int8_t action;

    setvbuf(stdout, NULL, _IONBF, 0);
    printf("Hello, World!\n");
    while (1) {
        if (1 != fread(&action, sizeof(action), 1, stdin))
            return 1;
        fprintf(stderr, "DEBUG: received action 0x%02x\r\n", action);

        switch (action) {
            case 'c': recv_variable(&c, &clen); break;
            case 'm': recv_variable(&m, &mlen); break;
            case 'a': recv_variable(&ad, &adlen); break;
            case 'k': recv_variable(&k, &klen); break;

            case 's': recv_variable(&nsec, &nslen); break;
            case 'p': recv_variable(&npub, &nplen); break;

            case 'C': send_variable(c, clen); break;
            case 'M': send_variable(m, mlen); break;
            case 'A': send_variable(ad, adlen); break;
            case 'K': send_variable(k, klen); break;
            case 'S': send_variable(nsec, nslen); break;
            case 'P': send_variable(npub, nplen); break;


            case 'e':
            case 'd':
/*
                fprintf(stderr, "mlen = %llu\r\n", mlen);
                fprintf(stderr, "m = "); FPRINTF_HEX(stderr, m, mlen); fprintf(stderr, "\r\n");
                fprintf(stderr, "adlen = %llu\r\n", adlen);
                fprintf(stderr, "ad = "); FPRINTF_HEX(stderr, ad, adlen); fprintf(stderr, "\r\n");
                fprintf(stderr, "nplen = %llu\r\n", nplen);
                fprintf(stderr, "np = "); FPRINTF_HEX(stderr, npub, nplen); fprintf(stderr, "\r\n");
                fprintf(stderr, "nslen = %llu\r\n", nslen);
                fprintf(stderr, "ns = "); FPRINTF_HEX(stderr, nsec, nslen); fprintf(stderr, "\r\n");
                fprintf(stderr, "klen = %llu\r\n", klen);
                fprintf(stderr, "k = "); FPRINTF_HEX(stderr, k, klen); fprintf(stderr, "\r\n");
*/

                if (k == NULL) {
                    fprintf(stderr, "Missing key\r\n");
                    return 3;
                }

                if (npub == NULL) {
                    fprintf(stderr, "Missing nonce\r\n");
                    return 3;
                }

                if (action == 'e') {

                    if (c == NULL) {
                        fprintf(stderr, "Missing ciphertext buffer\r\n");
                        return 3;
                    }

                    fprintf(stderr, "Starting encryption\r\n");
                    res = crypto_aead_encrypt(c, &clen,
                            m, mlen, ad, adlen, nsec, npub, k);
                    fprintf(stderr, "encryption finished %d\r\n", res);

                } else if (action == 'd') {

                    if (m == NULL) {
                        fprintf(stderr, "Missing message buffer\r\n");
                        return 3;
                    }

                    fprintf(stderr, "starting decryption\r\n");
                    res = crypto_aead_decrypt(m, &mlen,
                            nsec, c, clen, ad, adlen, npub, k);
                    fprintf(stderr, "decryption finished %d\r\n", res);

                }
                break;

            default:
                fprintf(stderr, "ERROR: unexpected action 0x%02x\r\n", action);
                return 2;
        }
        
    }

}