#include <zeta.h>
#include <tty.h>
#include <hardware.h>
#include <fifo.h>

#include <assert.h>
#include <stddef.h>
#include <string.h>

#ifdef __GNUC__
#define PACKED __attribute__((packed))
#else
#define PACKED
#endif

enum header_type {
    CMD_PING,
    CMD_INFO,
    CMD_BOOT,
    CMD_READ,
    CMD_WRITE,
    CMD_IO_READ,
    CMD_IO_WRITE,
    CMD_ECHO
};

struct header {
    uint8_t type;
    uint8_t bank;
    uint16_t address;
    uint16_t length;
    uint16_t checksum;
} PACKED;

// static_assert(sizeof(struct header) == 8, "struct header is not PACKED");

enum error {
    ERR_TIMEOUT = -1
};

#define MAX_PACKET_SIZE 256
#define TIMEOUT_MS 500
#define MAX_ATTEMPTS 3
#define MAX_TRANS_ATTEMPTS 5

uint16_t
crc16(const void *buf, size_t len);

volatile struct fifo rx_fifo = {0, 0, {0}};

enum ack {
    ACK = 0x00,
    NACK = 0xFF
};

void
rx_isr(void) __critical __interrupt(0)
{
    fifo_push(&rx_fifo, sio_a_data);
}

static volatile uint32_t millis = 0;

void
ctc1_isr(void) __critical __interrupt(1)
{
    millis += 5;
}

uint32_t
clock(void)
{
    volatile uint32_t ret;
    DI;
    ret = millis;
    EI;
    return ret;
}

void
putbyte(unsigned char b)
{
    unsigned char ctrl = 0;

    sio_a_data = b;

    while (!(ctrl & 0x04)) {
        sio_a_ctrl = 0;
        ctrl = sio_a_ctrl;
    }
}

static volatile int32_t errno = 0;

uint8_t
getbyte(void)
{
    uint8_t b;
    uint32_t ms = clock();
    errno = 0;
    while (fifo_empty(&rx_fifo)) {
        if (clock() - ms > TIMEOUT_MS) {
            errno = ERR_TIMEOUT;
            return 0;
        }
    }
    DI;
    b = fifo_pop(&rx_fifo);
    EI;
    return b;
}

void
flush(void)
{
    DI;
    fifo_clear(&rx_fifo);
    EI;
}


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

    for (uint8_t i = 0; i < 4; ++i)
        y ^= ((x >> i) & 1) ? c[i] : 0;

    return y;
}

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

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

    // Assume simple error, attempt correction
    if (p) {
        size_t i = 0;

        for (i = 0; i < LENGTH(r); ++i) {
            if (r[i] == x)
                break;
        }

        x ^= (1 << i);
    }

    return x & 0x0F;
}


int
read(void *buf, size_t count)
{
    uint8_t b;
    uint8_t *p = buf;

    for (int n = 0; n < count; ++n) {
        b = decode(getbyte());
        if (errno)
            return errno;

        b |= (decode(getbyte()) << 4);
        if (errno)
            return errno;

        p[n] = b;
    }

    return 0;
}

int
write(const void *buf, size_t count)
{
    const uint8_t *p = buf;

    for (size_t i = 0; i < count; ++i) {
        putbyte(encode(p[i] & 0x0F));
        putbyte(encode((p[i] >> 4) & 0x0F));
    }

    return 0;
}

int
read_header(struct header *header)
{
    int err;
    uint8_t ack;
    uint16_t checksum;

    while (1) {
        if ((err = read(header, sizeof(*header))))
            return err;

        checksum = header->checksum;
        header->checksum = 0;

        if (checksum == crc16(header, sizeof(*header))) {
            header->checksum = checksum;
            ack = ACK;
            write(&ack, sizeof(ack));
            return 0;
        } else {
            ack = NACK;
            write(&ack, sizeof(ack));
        }
    }
}

int
read_buf(size_t len, void *buf)
{
    int err;
    uint8_t ack;
    uint16_t checksum;

    for (int i = 0; i < MAX_TRANS_ATTEMPTS; ++i) {
        // TODO: reduce code?
        if ((err = read(&checksum, sizeof(checksum)))
            || (err = read(buf, len)))
            break;

        if (checksum == crc16(buf, len)) {
            ack = ACK;
            write(&ack, sizeof(ack));
            return 0;
        } else {
            ack = NACK;
            write(&ack, sizeof(ack));
        }
    }

    return -1;
}

int
write_buf(size_t len, const void *buf)
{
    int err;
    uint8_t ack = NACK;
    uint16_t checksum = crc16(buf, len);

    for (int i = 0; i < MAX_TRANS_ATTEMPTS; ++i){
        write(&checksum, sizeof(checksum));
        write(buf, len);

        // If TIMEOUT sending just give up
        if ((err = read(&ack, sizeof(ack))))
            return err;

        if (ack == ACK)
            return 0;
    }

    return -1;
}

// TODO: Restart after timeouts
void
bootloader(void)
{
    struct header header;
    uint8_t buf[MAX_PACKET_SIZE];

	const char *msg = "Bootloader ...\r\n";
	addstr(msg);

    while (1) {
        if (read_header(&header)) {
            flush();
            continue;
        }

        switch (header.type) {
        case CMD_BOOT:
            ((void (*)(void))header.address)();
            break;

        case CMD_READ:
            write_buf(header.length, (const void *)header.address);
            break;

        case CMD_WRITE:
            if (!read_buf(header.length, buf))
                memcpy((void *)header.address, buf, header.length);
            break;

        case CMD_ECHO:
            if (!read_buf(header.length, buf))
                write_buf(header.length, buf);
            break;

        default:
            break;
        }

        flush();
    }
}