#include "crypto_aead.h"
#include "api.h"

#define MAX_BYTES 100

#define DEBUG

uint8_t npub[CRYPTO_NPUBBYTES];
uint8_t nsec[CRYPTO_NSECBYTES];
uint8_t k[CRYPTO_KEYBYTES];
uint8_t ad[MAX_BYTES];
unsigned long long int adlen;
uint8_t m[MAX_BYTES];
unsigned long long int mlen;
uint8_t c[MAX_BYTES];
unsigned long long int clen;

void uart_write(uint8_t x) {
    Serial.write((uint8_t) (x));
}

uint8_t uart_read() {
    int r;
    do {
        r = Serial.read();
    } while (r == -1);
    return (uint8_t) (0xff & r);
}

// Simple serial protocol with packets and checksum
const uint8_t AMUX_TAG = 0xf9;
const uint8_t AMUX_END = 0xf3;
const uint8_t AMUX_EXT = 0x80;

void uartp_send(const void *src, uint16_t len) {
    uint8_t len_ind_0, len_ind_1, fcs, info;
    const uint8_t *buf = (const uint8_t *) src;

    uart_write(AMUX_TAG);
    len_ind_0 = (uint8_t) (0xff & len);
    len_ind_1 = (uint8_t) (0xff & (len >> 7));
    if (len < 128) {
        uart_write(len_ind_0);
    } else {
        uart_write(len_ind_0 | AMUX_EXT);
        uart_write(len_ind_1);
    }
    fcs = 0;
    for (uint16_t i = 0; i < len; i++) {
        info = buf[i];
        fcs += info;
        uart_write(buf[i]);
    }
    fcs = 255 - fcs;
    uart_write(fcs);
    uart_write(AMUX_END);
}


uint16_t uartp_recv(void *dst, uint16_t buf_len) {
    uint8_t *buf = (uint8_t *) dst;
    uint8_t tag_old, tag, info, cs;
    uint16_t len;

    tag = AMUX_END;
    while (1) {

        do {
            tag_old = tag;
            tag = uart_read();
        } while(tag != AMUX_TAG || tag_old != AMUX_END);

        len = (uint16_t) uart_read();
        if (len & AMUX_EXT) {
            len &= (~AMUX_EXT);
            len |= (uint16_t) (uart_read() << 7);
        }
        if (len > buf_len) {
            return len;
        }

        uint16_t i = 0;
        cs = 0;
        for (i = 0; i < len; i++) {
            info = uart_read();
            buf[i] = info;
            cs += info;
        }
        cs += uart_read();
        tag = uart_read();
        if (0xff == cs) {
            if (AMUX_END == tag) {
                return len;
            }
        }

    }
}

#ifdef DEBUG
#define DEBUG_BUF_LEN 80
int dbg_printf(const char *format, ...) {
    char printbuf[DEBUG_BUF_LEN+2];

    va_list vargs;
    va_start(vargs, format);
    int r = vsnprintf(printbuf+1, DEBUG_BUF_LEN, format, vargs);
    va_end(vargs);

    if (r < 0) {
        memcpy(printbuf+1, "DEBUG ERROR\r\n", 13);
        r = 13;
    }

    r = r > DEBUG_BUF_LEN ? DEBUG_BUF_LEN : r;
    printbuf[0] = 0xde;  // Debug messages should start with "\xde"
    uartp_send(printbuf, r+1);
    return r;
}
#else
  #define dbg_printf(...) (0)
#endif

void assert(bool b) {
    if (b)
        return;
    dbg_printf("Assertion failed\r\n");
    for(;;)
        yield();
}

void setup() {
    Serial.begin(115200);
    pinMode(LED_BUILTIN, OUTPUT);
    digitalWrite(LED_BUILTIN, HIGH);
    delay(100);
    Serial.println("Hello, World!");
}


void loop() {
    static uint8_t buf[256];
    int res;
    uint16_t len = uartp_recv(buf, 255);
    uint8_t action = buf[0];
    if (len == 0 || len > 255)
        return;

    dbg_printf("Received command %c\r\n", action);
    uint16_t l = len - 1;
    uint8_t *var = buf+1;
    switch (action) {
        case 'm': assert(l <= MAX_BYTES);        memcpy(m,    var, l); mlen = l; break;
        case 'c': assert(l <= MAX_BYTES);        memcpy(c,    var, l); clen = l; break;
        case 'a': assert(l <= MAX_BYTES);        memcpy(ad,   var, l); adlen = l; break;
        case 'k': assert(l == CRYPTO_KEYBYTES);  memcpy(k,    var, l); break;
        case 'p': assert(l == CRYPTO_NPUBBYTES); memcpy(npub, var, l); break;
        case 's': assert(l == CRYPTO_NSECBYTES); memcpy(nsec, var, l); break;
        case 'e':
            digitalWrite(LED_BUILTIN, LOW);
            res = crypto_aead_encrypt(c, &clen, m, mlen, ad, adlen, nsec, npub, k);
            digitalWrite(LED_BUILTIN, HIGH);
            break;
        case 'd':
            digitalWrite(LED_BUILTIN, LOW);
            res = crypto_aead_decrypt(m, &mlen, nsec, c, clen, ad, adlen, npub, k);
            digitalWrite(LED_BUILTIN, HIGH);
            break;
    }

    switch (action) {
        case'a': case'c': case'k': case'm': case'p': case's':
        case'd': case'e':
            l = 0;
            break;
        case'M': var = m;    l = mlen; break;
        case'C': var = c;    l = clen; break;
        case'A': var = ad;   l = adlen; break;
        case'K': var = k;    l = CRYPTO_KEYBYTES; break;
        case'P': var = npub; l = CRYPTO_NPUBBYTES; break;
        case'S': var = nsec; l = CRYPTO_NSECBYTES; break;
    }
    buf[0] = action;
    memcpy(buf+1, var, l);
    uartp_send(buf, l+1);
}