#define MODULE "ota"
#define DEBUG 1

#include "common.h"
#include "fw.h"

#include <esp_ota_ops.h>

#define BUFFER_SIZE	4096

esp_err_t esp_update(read_func_t read_func, token_t token, size_t size)
{
    const esp_partition_t *partition;
    esp_err_t err = ESP_ERR_INVALID_ARG;
    esp_ota_handle_t handle;
    bool ota_started = false;
    int len;
    char *buf = NULL;

    buf = malloc(BUFFER_SIZE);
    if (!buf) {
	err = ESP_ERR_NO_MEM;
	goto fail;
    }

    partition = esp_ota_get_next_update_partition(NULL);
    if (!partition) {
	err = ESP_ERR_NOT_FOUND;
	goto fail;
    }

    MSG("starting ESP OTA update, total %zu bytes\n", size);

    err = esp_ota_begin(partition, size, &handle);
    if (err)
	goto fail;

    ota_started = true;

    size_t left = size;
    while (left) {
	size_t chunk = left;
	if (chunk > BUFFER_SIZE)
	    chunk = BUFFER_SIZE;

	len = read_func(token, buf, chunk);
	if (len <= 0)
	    break;

	left -= len;

	MSG("writing %d bytes, %zu left\n", len, left);
	err = esp_ota_write(handle, buf, len);
	if (err)
	    goto fail;
    }

    if (left) {
	MSG("short input data, aborting update\n");
	err = ESP_ERR_OTA_VALIDATE_FAILED;
	goto fail;
    }

    MSG("finalizing flash update\n");
    err = esp_ota_end(handle);
    if (err)
	goto fail;

    free(buf);

    MSG("setting boot partition\n");
    return esp_ota_set_boot_partition(partition);

fail:
    MSG("failure: error 0x%x\n", err);

    if (ota_started)
	esp_ota_abort(handle);
    if (buf)
	free(buf);
    return err;
}