//
// Created by nock on 26/03/19.
//

#include <errno.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <netinet/in.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <math.h>

#include "base64.h"
#include "crc.h"
#include "dump.h"
#include "error.h"
#include "fixkey-msg.h"
#include "fixkey-payload.h"
#include "gpio.h"
#include "image.h"
#include "nmgr.h"
#include "timer.h"
#include "serial_util.h"
#include <tinycbor/cbor.h>

#ifdef DO_UPLOAD_SHA
#include "sha256.h"
#endif


/* These values produce output identical to mcumgr *if* DO_UPLOAD_SHA is defined */
#define IMAGE_CHUNK_SIZE 124 // 124 for identical output to newtmgr
#define MACRO_CHUNK_SIZE 302 // 302 ""

/* Fast mode doesn't support original EM bootloader, but is much faster with later bootloaders */
#define IMAGE_CHUNK_SIZE_FAST 216
#define MACRO_CHUNK_SIZE_FAST 432

/*
 * Larger values above greatly increase speed, but seem to break older
 * versions of mcuboot (like the one shipped with inital EM firmware)
 */

#define CRC16_INITIAL_CRC       0       /* what to seed crc16 with */
#define CRC_CITT_POLYMINAL 0x1021

static int serial_fd = -1;

static bool g_fast_mode = false;
static bool g_no_fix = false;

void image_set_fast_mode(void) {
    g_fast_mode = true;
}

static bool image_get_fast_mode(void) {
    return g_fast_mode;
}

void image_help(void){
    BAIL_OUT(ERR_USAGE);
}

static void read_response(uint8_t *dest, size_t *len) {
    static uint8_t buf[BOOT_SERIAL_IN_MAX*2] = {0};
    static uint8_t *p = buf;
    timer_start(3);
    while (true) {
        int bytes_read = read(serial_fd, p, BOOT_SERIAL_IN_MAX);
        if (bytes_read < 0) {
            if (errno == EAGAIN) {
                if (timer_expired()) {
                    dump_raw(buf, p-buf, DUMP_RXI);
                    BAIL_OUT(ERR_TIMEOUT);
                }
                usleep(5000);
                continue;
            }
            BAIL_OUT(ERR_SERIAL_PORT);
        }
        dump_raw(p, bytes_read, DUMP_RXI_BUFFER_UPDATE);
        p += bytes_read;
        uint8_t *nl_p = memchr(buf, '\n', p-buf);
        if (nl_p) {
            *len = nl_p - buf;
            memcpy(dest, buf, *len);
            uint16_t new_buf_len = p - nl_p;
            if (new_buf_len) {
                p = buf;
                return;
            }
            memcpy(buf, nl_p+1, new_buf_len);
            p = buf+new_buf_len;
            break;
        }
    }
}



static bool assert_pkt_start(uint8_t const *buf){
    return (buf[0] == SHELL_NLIP_PKT_START1 && buf[1] == SHELL_NLIP_PKT_START2);
}

#define cbor_require_type(it, type) _cbor_require_type(it, type, __FUNCTION__, __FILE__, __LINE__)
static void _cbor_require_type(CborValue *it, CborType type, const char *function, const char *file, unsigned line) {
    CborType actual_type = cbor_value_get_type(it);
    if (actual_type != type) {
        _bail_out(ERR_CBOR_TYPE, function, file, line);
    }
}

#define cbor_ensure_success(cbor_err) _cbor_ensure_success(cbor_err, __FUNCTION__, __FILE__, __LINE__)
static inline void _cbor_ensure_success(CborError err, const char *function, const char *file, unsigned line) {
    if (err != CborNoError) {
        _bail_out(ERR_CBOR, function, file, line);
    }
}

static char *get_version(void) {
    serial_flush(serial_fd);
    int rc = write(serial_fd, IMAGE_LIST_CMD, sizeof(IMAGE_LIST_CMD) - 1);
    dump_raw((uint8_t *)IMAGE_LIST_CMD, sizeof(IMAGE_LIST_CMD)-1, DUMP_TXO);
    if (rc < sizeof(IMAGE_LIST_CMD) - 1) {
        BAIL_OUT(ERR_SERIAL_PORT);
    }
    uint8_t tmp[BOOT_SERIAL_IN_MAX];
    size_t in_len;

    read_response(tmp, &in_len);
    dump_raw(tmp, in_len, DUMP_RXI);

    // Process the response
    CborParser parser;
    CborValue it;

    // Validate packet start marker
    if (in_len < 2) {
        BAIL_OUT(ERR_INVALID_PACKET);
    }

    if (!assert_pkt_start(tmp)) {
        BAIL_OUT(ERR_INVALID_PACKET);
    }

    // Decode the payload
    size_t out_len;
    uint8_t *decoded_buf = base64_decode(tmp+2, in_len-2, &out_len);

    if (!decoded_buf) {
        BAIL_OUT(ERR_BASE64_DECODE);
    }

    // Unclear why this would be useful in a line oriented (cannonical mode) program
    uint16_t pkt_len = ntohs(*(uint16_t *)decoded_buf);
    uint16_t crc = crc16(decoded_buf+sizeof(pkt_len), pkt_len, CRC_CITT_POLYMINAL, CRC16_INITIAL_CRC, true);
    if (crc) {
        BAIL_OUT(ERR_BAD_CRC);
    }

    uint8_t *p = decoded_buf + sizeof(pkt_len) + sizeof(struct nmgr_hdr);
    cbor_ensure_success(
        cbor_parser_init(p, out_len, 0, &parser, &it));
    cbor_require_type(&it, CborMapType);
    CborValue map;
    cbor_ensure_success(cbor_value_enter_container(&it, &map));
    cbor_require_type(&map, CborTextStringType);
    cbor_ensure_success(cbor_value_advance(&map));
    cbor_require_type(&map, CborArrayType);
    CborValue img_array;
    cbor_ensure_success(cbor_value_enter_container(&map, &img_array));

    char *version = NULL;
    while (!cbor_value_at_end(&img_array)) {
        cbor_require_type(&img_array, CborMapType);
        CborValue version_tag;
        cbor_ensure_success(
            cbor_value_map_find_value(&img_array, "version", &version_tag));

        if (cbor_value_get_type(&version_tag) == CborInvalidType) {
            cbor_value_advance(&img_array);
            continue;
        }
        size_t len;
        cbor_value_dup_text_string(&version_tag, &version, &len, NULL);
        break;
    }
    free(decoded_buf);
    return version;
}

static void wrap_and_send_pkt(uint8_t *data, size_t len, bool fast_mode) {
    assert(data != NULL);
    assert(len != 0);
    static uint8_t seq = 67;
    uint8_t b64_buf[BOOT_SERIAL_OUT_MAX*2];
    uint16_t tot_len = 2 + len + sizeof(struct nmgr_hdr) + 2; // +2 is len +2 is crc
    uint16_t net_len = htons(tot_len-2); // tot_len doesn't include the len field

    struct nmgr_hdr hdr = {0};
    hdr.nh_op = NMGR_OP_WRITE;
    hdr.nh_group = htons(MGMT_GROUP_ID_IMAGE);
    hdr.nh_id = IMGMGR_NMGR_ID_UPLOAD;
    hdr.nh_seq = seq++;
    hdr.nh_len = htons(len);

    memcpy(b64_buf, &net_len, sizeof(net_len));

#define HDR_OFFSET 2
#define DATA_OFFSET (HDR_OFFSET+sizeof(struct nmgr_hdr))
#define CRC_OFFSET(len) (DATA_OFFSET+len)

    memcpy(b64_buf+HDR_OFFSET, &hdr, sizeof(struct nmgr_hdr));
    memcpy(b64_buf+DATA_OFFSET, data, len);

    uint16_t crc = crc16(b64_buf+HDR_OFFSET, tot_len-4, CRC_CITT_POLYMINAL, CRC16_INITIAL_CRC, true);
    crc = htons(crc);
    memcpy(b64_buf+CRC_OFFSET(len), &crc, sizeof(crc));

    size_t out_len;
    uint8_t *out_buf = base64_encode(b64_buf, tot_len, &out_len);

    uint8_t *pos = out_buf;
    while (pos < out_buf+out_len) {
        uint8_t start_pkt[] = {SHELL_NLIP_PKT_START1, SHELL_NLIP_PKT_START2};
        uint8_t start_data[] = {SHELL_NLIP_DATA_START1, SHELL_NLIP_DATA_START2};

        if (pos == out_buf) {
            write(serial_fd, start_pkt, sizeof(start_pkt));
            dump_raw(start_pkt, sizeof(start_pkt), DUMP_TXO);
        } else {
            if (!fast_mode) {
                usleep(20000);
            }
            write(serial_fd, start_data, sizeof(start_data));
            dump_raw(start_data, sizeof(start_data), DUMP_TXO);
        }

        size_t remaining = out_buf + out_len - pos;
        size_t to_write;
        if (fast_mode) {
            to_write = remaining > IMAGE_CHUNK_SIZE_FAST ? IMAGE_CHUNK_SIZE_FAST : remaining;
        } else {
            to_write = remaining > IMAGE_CHUNK_SIZE ? IMAGE_CHUNK_SIZE : remaining;
        }

        write(serial_fd, pos, to_write);
        dump_raw(pos, to_write, DUMP_TXO);
        write(serial_fd, "\n", 1);
        dump_raw("\n", 1, DUMP_TXO);
        pos += to_write;
    }
    free(out_buf);
    fsync(serial_fd);
}

void do_upload(char *filename, bool fast_mode){
    serial_flush(serial_fd);
    size_t pos = 0, len;
    uint8_t const *payload;
    if (filename) {
        int fd = open(filename, O_RDONLY);
        if (fd < 0) {
            fprintf(stderr, "Failed to open \"%s\": %s\n", filename, strerror(errno));
            exit(ERR_SERIAL_PORT);
        }

        struct stat st;
        fstat(fd, &st);
        len = st.st_size;

        uint8_t *payload_mut = calloc(1, len);
        if (!payload_mut) {
            BAIL_OUT(ERR_MEMORY);
        }

        size_t bytes_read = 0;
        while (bytes_read < len) {
            errno = 0;
            int r = read(fd, &payload_mut[bytes_read], len);
            if (r < 0 && errno != EAGAIN) {
                BAIL_OUT(ERR_IO);
            }
            bytes_read += r;
        }
        payload = payload_mut;
    } else {
        payload = fixkey_payload_bin;
        len = fixkey_payload_len;
    }

    size_t base_chunk_len;
    if (fast_mode) {
        base_chunk_len = MACRO_CHUNK_SIZE_FAST;
    } else {
        base_chunk_len = MACRO_CHUNK_SIZE;
    }
    while (pos < len){
        size_t data_chunk_len;
        if (pos == 0) {
            data_chunk_len = base_chunk_len;
        } else {
            data_chunk_len = base_chunk_len + 45;
        }
        CborEncoder root, map;
        uint8_t cbor_buf[BOOT_SERIAL_OUT_MAX*2];
        cbor_encoder_init(&root, cbor_buf, BOOT_SERIAL_OUT_MAX*2, 0);
#ifdef DO_UPLOAD_SHA
        cbor_ensure_success(cbor_encoder_create_map(&root, &map, pos == 0 ? 4 : 2));
#else
        cbor_ensure_success(cbor_encoder_create_map(&root, &map, pos == 0 ? 3 : 2));
#endif

        cbor_ensure_success(cbor_encode_text_stringz(&map, "data"));

        // Most packets will be `data_chunk_len`, but not the last one
        size_t actual_size = len - pos > data_chunk_len ? data_chunk_len : len - pos;

        cbor_ensure_success(cbor_encode_byte_string(&map, &payload[pos], actual_size));

        if (pos == 0) {
            cbor_ensure_success(cbor_encode_text_stringz(&map, "len"));
            cbor_ensure_success(cbor_encode_uint(&map, len));
        }

        cbor_ensure_success(cbor_encode_text_stringz(&map, "off"));
        cbor_ensure_success(cbor_encode_uint(&map, pos));

#ifdef DO_UPLOAD_SHA
        if (pos == 0) {
            cbor_ensure_success(cbor_encode_text_stringz(&map, "sha"));
            SHA256_CTX sha_ctx;
            sha256_init(&sha_ctx);
            sha256_update(&sha_ctx, payload, len);
            uint8_t sha[32];
            sha256_final(&sha_ctx, (BYTE *)&sha);
            cbor_ensure_success(cbor_encode_byte_string(&map, sha, 32));
        }
#endif /* DO_UPLOAD_SHA */

        cbor_ensure_success(cbor_encoder_close_container(&root, &map));
        size_t cbor_size = cbor_encoder_get_buffer_size(&root, cbor_buf);

        wrap_and_send_pkt(cbor_buf, cbor_size, fast_mode);

        uint8_t resp_buf[BOOT_SERIAL_IN_MAX*2];
        size_t in_len;
        read_response(resp_buf, &in_len);

        size_t out_len;
        uint8_t *out = base64_decode(resp_buf+2, in_len-2, &out_len);

        if (out == NULL) {
            BAIL_OUT(ERR_BASE64_DECODE);
        }

        uint16_t l;
        memcpy(&l, out, sizeof(l));
        l = ntohs(l);

        if (crc16(out+2, l, CRC_CITT_POLYMINAL, CRC16_INITIAL_CRC, true)) {
            BAIL_OUT(ERR_BAD_CRC);
        }

        CborParser parser;
        CborValue it, rc;
        cbor_ensure_success(
            cbor_parser_init(out+sizeof(struct nmgr_hdr)+2, out_len-sizeof(struct nmgr_hdr)-2, 0, &parser, &it));
        cbor_require_type(&it, CborMapType);
        if (cbor_value_map_find_value(&it, "rc", &rc)) {
            BAIL_OUT(ERR_NO_RESPONSE_CODE);
        }
        int r;
        cbor_ensure_success(cbor_value_get_int(&rc, &r));
        if (r) {
            BAIL_OUT(ERR_CLIENT_REPORTED);
        }
        if (cbor_value_map_find_value(&it, "off", &rc)) {
            fprintf(stderr, "Failed to find an offset in the response %s:%d\n", __FILE__, __LINE__);
        } else {
            cbor_ensure_success(cbor_value_get_int(&rc, &r));
            uint8_t percent_complete = round(((double)r/len)*100);
            printf("\033[80D%d/%lu Written (%d%%)", r, len, percent_complete);
            fflush(stdout);
        }
        pos += actual_size;
        free(out);
    }
    printf("\n");
}

void fix_and_upload(char *filename) {
    char *version = get_version();
    bool has_new_bootloader = false;
    if (version && !strcmp(version, "0.3.2.0") ) {
        // Upload bootloader update payload
        printf("Bootloader update required, uploading...\n");
        do_upload(NULL, false);

        // Reset into controller
        printf("Starting Controller\n");
        gpio_ble_req_controller();
        sleep(1);

        // Update bootloader
        printf("Triggering bootloader upgrade\n");
        write(serial_fd, fixkey_msg, sizeof(fixkey_msg));
        sleep(3);

        // Reset into new bootloader
        printf("Reseting into new bootloader\n");
        gpio_ble_req_bootloader();
        usleep(50000);
        has_new_bootloader = true;
    }

    // Do the update
    printf("Writing %s into slot 1\n", filename);
    do_upload(filename, image_get_fast_mode() || has_new_bootloader); // If we just updated the BL, it supports fast
}

void image_set_no_fix(void) {
    g_no_fix = true;
}

void image_main(int fd, int argc, char **argv){
    gpio_ble_req_bootloader();
    serial_fd = fd;
    if (argc < 1) {
        image_help();
        return;
    }
    if (!strcmp(argv[0], "version")) {
        char *version = get_version();

        if (version) {
            printf("%s\n", version);
            free(version);
        } else {
            printf("No application image found.\n");
        }
    } else if (!strcmp(argv[0], "upload")) {
        if (argc < 2) {
            fprintf(stderr, "umcumgr image upload <binary file>\n");
        }
        if (g_no_fix) {
            printf("Writing %s into slot 1\n", argv[1]);
            do_upload(argv[1], image_get_fast_mode());
        } else {
            fix_and_upload(argv[1]);
        }
    } else {
        image_help();
    }
    gpio_ble_req_controller();
}