#include <stdio.h>
#include <fcntl.h>
#include <string.h>
#include <unistd.h>
#include "system.h"
#include "spi_flash.h"
#include "wolfboot/include/target.h"

extern void wolfBoot_success(void);
#define MSGSIZE (4 + 4 + 8)
#define PAGESIZE (256)

uint8_t page[PAGESIZE];

static uint32_t next_seq;
static uint32_t tot_len = 0;
const char err[]="!!!!";


static void ack(uint32_t _off)
{
    uint32_t off = _off;
    uint8_t ack_start = '#';
    
    write(STDOUT_FILENO, &ack_start, 1);
    write(STDOUT_FILENO, &off, 4);
}

static int check(uint8_t *pkt, int size)
{
    int i;
    uint16_t c = 0;
    uint16_t c_rx = *((uint16_t *)(pkt + 2));
    uint16_t *p = (uint16_t *)(pkt + 4);
    for (i = 0; i < ((size - 4) >> 1); i++)
        c += p[i];
    if (c == c_rx)
        return 0;
    return -1;
}

static uint8_t msg[MSGSIZE];

static void update_loop(void)
{
    int r;
    uint32_t tlen = 0;
    volatile uint32_t recv_seq;
    uint32_t r_total = 0;
    memset(page, 0xFF, PAGESIZE);
    while (1) {
        r_total = 0;
        do {
            while(r_total < 2) {
                r = read(STDIN_FILENO, msg + r_total, 2 - r_total);
                if (r <= 0)
                    continue;
                r_total += r;
                if ((r_total == 2) && ((msg[0] != 0xA5) || msg[1] != 0x5A)) {
                    r_total = 0;
                    continue;
                }
            }
            r = read(STDIN_FILENO, msg + r_total, MSGSIZE - r_total);
            if (r <= 0)
                continue;
            r_total += r;
            if ((tot_len == 0) && r_total == 2 + sizeof(uint32_t))
                break;
            if ((r_total > 8)  && (tot_len <= ((r_total - 8) + next_seq)))
                break;
        } while (r_total < MSGSIZE);

        if (tot_len == 0)  {
            tlen = msg[2] + (msg[3] << 8) + (msg[4] << 16) + (msg[5] << 24);
            if (tlen > WOLFBOOT_PARTITION_SIZE - 8) {
                write(STDOUT_FILENO, err, 4);
                break;
            }
            tot_len = tlen;
            ack(0);
            continue;
        }

        if (check(msg, r_total) < 0) {
            ack(next_seq);
            continue;
        }

        recv_seq = msg[4] + (msg[5] << 8) + (msg[6] << 16) + (msg[7] << 24);
        if (recv_seq == next_seq)
        {
            int psize = r_total - 8;
            int page_idx = recv_seq % PAGESIZE;
            memcpy(&page[recv_seq % PAGESIZE], msg + 8, psize);
            page_idx += psize;
            if ((page_idx == PAGESIZE) || (next_seq + psize >= tot_len)) {
                uint32_t dst = (WOLFBOOT_PARTITION_UPDATE_ADDRESS + recv_seq) / PAGESIZE;
                flashpage_write_and_verify(dst, page);
                memset(page, 0xFF, PAGESIZE);
            }
            next_seq += psize;
        }
        ack(next_seq);
        if (next_seq >= tot_len) {
            /* Update complete */
            wolfBoot_update_trigger();
            reboot();
            while(1) 
                ;
        }
    }
}