#include <zup.h>

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

#include <fcntl.h>
#include <unistd.h>


// Hamming(7,4) encoding
static unsigned char
encode(unsigned char x)
{
    unsigned char y = 0;
    const unsigned char c[4] = {0x61, 0x52, 0x34, 0x78};

    for (int i = 0; i < 4; ++i)
        y ^= ((x >> i) & 1) * c[i];

    return y;
}

// Hamming(7,4) decoding
static unsigned char
decode(unsigned char x)
{
    unsigned char p = 0;
    const unsigned char r[7] = {6, 5, 3, 7, 1, 2, 4};

    for (int i = 0; i < 7; ++i)
        p ^= ((x >> i) & 1) * r[i];

    // Attempt correcting simple error
    if (p) {
        size_t n = 0;
        for (size_t i = 0; i < n; ++i) {
            if (r[i] == x)
                n = i;
        }
        x ^= (1 << n);
    }

    return x & 0x0F;
}

static void
encode_buf(size_t slen, const void *_src, size_t dlen, void *_dest)
{
    const uint8_t *src = _src;
    uint8_t *dest = _dest;

    for (size_t i = 0; i < slen && i * 2 < dlen; ++i) {
        dest[2 * i] = encode(src[i] & 0x0F);
        dest[2 * i + 1] = encode((src[i] >> 4) & 0x0F);
    }
}

static void
decode_buf(size_t slen, const void *_src, size_t dlen, void *_dest)
{
    const uint8_t *src = _src;
    uint8_t *dest = _dest;

    for (size_t i = 0; i < dlen && i * 2 < slen; ++i) {
        dest[i] = decode(src[2 * i]);
        dest[i] |= (decode(src[2 * i + 1]) << 4);
    }
}

// Ensure count bytes are read unless read times out MAX_ATTEMPTS times
static int
e_read(int fd, void *buf, size_t count)
{
    ssize_t ret;
    ssize_t nread = 0;
    uint8_t *p = buf;

    for (int n = 0; n < MAX_ATTEMPTS && nread < count; ++n) {
        while ((ret = read(fd, &p[nread], count - nread)) != 0) {
            if (ret < 0)
                goto error;
            else
                nread += ret;
        }
    }

    return (count == nread) ? 0 : ERR_TIMEOUT;
error:
    perror("e_read failed");
    return -1;
}

// Ensure count bytes are written
static int
e_write(int fd, const void *buf, size_t count)
{
    ssize_t ret = 0;
    ssize_t nwrite = 0;
    const uint8_t *p = buf;

    while (nwrite < count) {
        ret = write(fd, &p[nwrite], count - nwrite);
        if (ret < 0)
            goto error;
        nwrite += ret;
    }

    return 0;
error:
    perror("e_write failed!");
    return ERR_WRITE;
}

static int
read_ack(int fd)
{
    int err;
    uint8_t ack;
    uint8_t buf[2];
    if ((err = e_read(fd, buf, sizeof(buf)))) {
        msg(1, "read_ack failed\n");
        return err;
    }
    decode_buf(sizeof(buf), buf, sizeof(ack), &ack);
    msg(1, "read_ack success\n");
    return ack;
}

static int
write_ack(int fd, uint8_t ack)
{
    int err;
    uint8_t buf[2];
    encode_buf(sizeof(ack), &ack, sizeof(buf), buf);
    if ((err = e_write(fd, buf, sizeof(buf)))) {
        msg(1, "write_ack failed\n");
        return err;
    }
    tcdrain(fd);
    msg(1, "write_ack success\n");
    return 0;
}

static int
write_header(int fd, const struct header *header)
{
    int err;
    uint8_t buf[2 * sizeof(*header)];
    encode_buf(sizeof(*header), header, sizeof(buf), buf);

    msg(1, "write_header\n");

    for (int i = 0; i < param.max_attempts; ++i) {
        msg(1, "  attempt [%d/%d]\n", i + 1, param.max_attempts);

        if ((err = e_write(fd, buf, sizeof(buf))))
            goto error;
        tcdrain(fd);

        err = read_ack(fd);

        if (err == ACK) {
            msg(1, "  write_header success\n");
            return 0;
        }

        if (err != NACK)
            goto error;
    }

    err = ERR_NACK;

error:
    msg(1, "  write_header failed\n");
    return err;
}

static int
read_buf(int fd, size_t len, void *_buf)
{
    int err;
    uint8_t *buf = _buf;
    uint8_t ebuf[2 * MAX_PACKET_SIZE];

    uint16_t checksum = 0;
    uint8_t echecksum[4];

    msg(1, "read_buf\n");

    for (int i = 0; i < param.max_attempts; ++i) {
        msg(1, "  attempt [%d/%d]\n", i + 1, param.max_attempts);

        if ((err = e_read(fd, echecksum, sizeof(echecksum)))
            || (err = e_read(fd, ebuf, 2 * len)))
            goto error;

        decode_buf(sizeof(echecksum), echecksum, sizeof(checksum), &checksum);
        decode_buf(2 * len, ebuf, len, buf);

        if (checksum == crc16(buf, len)) {
            write_ack(fd, ACK);
            msg(1, "  read_buf success\n");
            return 0;
        } else {
            write_ack(fd, NACK);
        }
    }

    err = ERR_NACK;

error:
    msg(1, "  read_buf failed\n");
    return err;
}

static int
write_buf(int fd, size_t len, const void *_buf)
{
    int err;
    const uint8_t *buf = _buf;
    uint8_t ebuf[2 * MAX_PACKET_SIZE];

    uint16_t checksum = crc16(buf, len);
    uint8_t echecksum[4];

    encode_buf(sizeof(checksum), &checksum, LEN(echecksum), echecksum);
    encode_buf(len, buf, sizeof(ebuf), ebuf);

    msg(1, "write_buf\n");

    int ack;
    for (int i = 0; i < param.max_attempts; ++i) {
        msg(1, "  attempt [%d/%d]\n", i + 1, param.max_attempts);

        if ((err = e_write(fd, echecksum, LEN(echecksum)))
            || (err = e_write(fd, ebuf, 2 * len)))
            goto error;

        tcdrain(fd);

        err = read_ack(fd);

        if (err == ACK) {
            msg(1, "  write_buf success\n");
            return 0;
        }

        if (err != NACK)
            goto error;
    }

error:
    msg(1, "  write_buf failed\n");
    return ack < 0 ? ack : ERR_NACK;
}

// Sends a boot command
// returns 0 on success
int
z_boot(int fd, uint8_t bank, uint16_t address)
{
    struct header header = {
        .type = CMD_BOOT,
        .bank = bank,
        .address = address,
        .length = 0,
        .checksum = 0
    };

    header.checksum = crc16(&header, sizeof(header));

    int err;
    tcflush(fd, TCIOFLUSH);
    if ((err = write_header(fd, &header)))
        tcflush(fd, TCIOFLUSH);

    return err;
}

// Sends a single CMD_READ header and reads its response
// Does no splitting
// returns 0 on success
int
z_read(int fd, uint8_t bank, uint16_t address, uint16_t length, uint8_t *buf)
{
    assert(length <= MAX_PACKET_SIZE);

    struct header header = {
        .type = CMD_READ,
        .bank = bank,
        .address = address,
        .length = length,
        .checksum = 0
    };

    header.checksum = crc16(&header, sizeof(header));

    int err;
    tcflush(fd, TCIOFLUSH);
    if ((err = write_header(fd, &header))
        || (err = read_buf(fd, length, buf))) {
        tcflush(fd, TCIOFLUSH);
        return err;
    }

    return 0;
}

// Sends a single CMD_WRITE header and packet
// Does no splitting
// returns 0 on success
int
z_write(int fd, uint8_t bank, uint16_t address, uint16_t length,
        const uint8_t *buf)
{
    assert(length <= MAX_PACKET_SIZE);

    struct header header = {
        .type = CMD_WRITE,
        .bank = bank,
        .address = address,
        .length = length,
        .checksum = 0
    };

    header.checksum = crc16(&header, sizeof(header));

    int err;
    tcflush(fd, TCIOFLUSH);
    if ((err = write_header(fd, &header))
        || (err = write_buf(fd, length, buf))) {
        tcflush(fd, TCIOFLUSH);
        return err;
    }

    return 0;
}

// Sends a single CMD_ECHO header, packet and reads the response
// Does no splitting
// returns 0 on success
int
z_echo(int fd, uint16_t length, uint8_t *buf)
{
    assert(length <= MAX_PACKET_SIZE);

    struct header header = {
        .type = CMD_ECHO,
        .bank = 0,
        .address = 0,
        .length = length,
        .checksum = 0
    };

    header.checksum = crc16(&header, sizeof(header));

    int err = 0;
    tcflush(fd, TCIOFLUSH);
    if ((err = write_header(fd, &header))
        || (err = write_buf(fd, length, buf))
        || (err = read_buf(fd, length, buf))) {
        tcflush(fd, TCIOFLUSH);
        return err;
    }

    return 0;
}

int
open_tty(const char *port, int baud)
{
    int fd = open(port, O_RDWR);

    if (fd < 0) {
        fprintf(stderr, "Error: File '%s' could not be opened\n", port);
        exit(EXIT_FAILURE);
    }

    struct termios term;
    tcgetattr(fd, &term);

    cfmakeraw(&term);
    // read with 1s timeout
    term.c_cc[VTIME] = 5;
    term.c_cc[VMIN] = 0;
    cfsetspeed(&term, baud);

    tcsetattr(fd, TCSANOW, &term);
    tcflush(fd, TCIOFLUSH);

    return fd;
}