#define MODULE "ota"
#define DEBUG 1

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

#include <esp_ota_ops.h>

#define BUFFER_SIZE	4096

static int esp_err_to_fwupdate(int err)
{
    switch (err) {
    case ESP_OK:
	return Z_OK;
    case ESP_ERR_NO_MEM:
	return Z_MEM_ERROR;
    case ESP_ERR_NOT_FOUND:
	return FWUPDATE_ERR_ESP_NO_PARTITION;
    case ESP_ERR_INVALID_ARG:
    case ESP_ERR_OTA_PARTITION_CONFLICT:
    case ESP_ERR_OTA_SELECT_INFO_INVALID:
    case ESP_ERR_OTA_ROLLBACK_INVALID_STATE:
	return FWUPDATE_ERR_ESP_BAD_OTA;
    case ESP_ERR_FLASH_OP_TIMEOUT:
    case ESP_ERR_FLASH_OP_FAIL:
	return FWUPDATE_ERR_ESP_FLASH_FAILED;
    case ESP_ERR_OTA_VALIDATE_FAILED:
	return FWUPDATE_ERR_ESP_BAD_DATA;
    default:
	return FWUPDATE_ERR_UNKNOWN;
    }
}

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

    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\r", len, left);
	err = esp_ota_write(handle, buf, len);
	CMSG("\n");
	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);
    buf = NULL;

    MSG("setting boot partition\n");
    err = esp_ota_set_boot_partition(partition);

fail:
    err = esp_err_to_fwupdate(err);

    if (err) {
	MSG("failure: error %d\n", err);
	if (ota_started)
	    esp_ota_abort(handle);
    }
    if (buf)
	free(buf);

    return err;
}