//
// Created by jedi on 25.06.21.
//

#include "system.h"
#include "crc32.h"
#include "log.h"

#include <FreeRTOS.h>
#include <sysparam.h>
#include <spiflash.h>

#include <espressif/esp_system.h>
#include <rboot/rboot.h>
#include <rboot-ota/rboot-api.h>
#include <string.h>

#define min(a, b) \
 ({ __typeof__ (a) _a = (a); \
 __typeof__ (b) _b = (b); \
 _a < _b ? _a : _b; })

void system_clear_config() {
    vPortEnterCritical();
    uint32_t num_sectors = 0x2000 / sdk_flashchip.sector_size;
    uint32_t start = 0x00100000;
    for (uint32_t i = 0; i < num_sectors; i++) {
        spiflash_erase_sector(start + i * sdk_flashchip.sector_size);
    }
    if(sysparam_create_area(start, num_sectors, true) == SYSPARAM_OK) {
        sysparam_init(start, 0);
    }
    sysparam_init(start, start + 0x2000);
    sdk_system_restart();
}

void system_init_config() {
    uint32_t base_addr = 0x00100000;
    uint32_t num_sectors;
    sysparam_init(base_addr, 0);
    if(sysparam_get_info(&base_addr, &num_sectors) != SYSPARAM_OK) {
        syslog("Warning: WiFi config, sysparam not initialized\n");
        num_sectors = 0x2000 / sdk_flashchip.sector_size;
        if(sysparam_create_area(base_addr, num_sectors, true) == SYSPARAM_OK) {
            sysparam_init(base_addr, 0);
        }
        sdk_system_restart();
    }
}

#define MAX_IMAGE_SIZE 0x100000

struct {
    rboot_write_status status;
    uint32_t head;
    uint32_t base;
    uint16_t seq;
    uint8_t slot;
} otaflash_context;

void system_otaflash_init() {
    rboot_config conf;
    conf = rboot_get_config();
    otaflash_context.slot = (conf.current_rom + 1) % conf.count;
    otaflash_context.base = rboot_get_slot_offset(otaflash_context.slot);
    otaflash_context.status = rboot_write_init(otaflash_context.base);
    otaflash_context.head = otaflash_context.base;
    otaflash_context.seq = 0;
}

enum return_code system_otaflash_chunk(uint8_t *data, uint16_t len, uint16_t seq, uint32_t hash, uint16_t *ack) {
    uint32_t local_hash = crc32(data, len);
    if(hash == local_hash && otaflash_context.seq == seq) {
        if(otaflash_context.head % SECTOR_SIZE == 0) {
            sdk_spi_flash_erase_sector(otaflash_context.head / SECTOR_SIZE);
        }
        if(((uint32_t) data) % 4) {
            uint32 buf[len / 4];
            memcpy(buf, data, len);
            sdk_spi_flash_write(otaflash_context.head, buf, len);
        } else {
            sdk_spi_flash_write(otaflash_context.head, (uint32_t *) data, len);
        }
        otaflash_context.head += len;
        otaflash_context.seq++;
        return OK;
    } else if(hash != local_hash) {
        return CHECKSUM_MISMATCH;
    } else {
        if(ack)
            *ack = otaflash_context.seq;
        return SEQUENCE_OUT_OF_ORDER;
    }

}

void system_otaflash_verify_chunk(void *ctx, void *data, size_t len) {
    uint32_t digest = *(uint32_t *) ctx;
    digest = crc32_partial(digest, data, len);
    *(uint32_t *) ctx = digest;
}

enum return_code system_otaflash_verify_and_switch(uint32_t len, uint32_t hash) {

    uint32_t digest = 0;
    rboot_digest_image(otaflash_context.base, min(len, MAX_IMAGE_SIZE), system_otaflash_verify_chunk, &digest);

    if(hash != digest) {
        syslog("OTA failed to verify firmware\r\n");
        return CHECKSUM_MISMATCH;
    }

    vPortEnterCritical();
    if(!rboot_set_current_rom(otaflash_context.slot)) {
        syslog("OTA Update failed to set new rboot slot\r\n");
        vPortExitCritical();
        return RBOOT_SWITCH_FAILED;
    }
    vPortExitCritical();
    return OK;
}