#define MODULE "fwupdate"
#define DEBUG 1

#include "common.h"
#include "jtag.h"
#include "spiflash.h"
#include "fpga.h"
#include "ota.h"
#include "spz.h"
#include "httpd.h"
#include "fw.h"
#include "boardinfo_esp.h"

#include <unzipLIB.h>
#include <zlib.h>

/* Needed for struct inflate_state, due to unziplib hacks */
#include <zutil.h>
#include <inftrees.h>
#include <inflate.h>
#ifndef local
# define local static
#endif

#define BUFFER_SIZE		SPIFLASH_SECTOR_SIZE
#define FWUPDATE_STACK		8192
#define FWUPDATE_PRIORITY	3

static void heap_info(void)
{
#if DEBUG > 1
    MSG("Heap: sram ");
    MSG("%u/", heap_caps_get_largest_free_block(MALLOC_CAP_INTERNAL));
    MSG("%u, spiram ", heap_caps_get_free_size(MALLOC_CAP_INTERNAL));
    MSG("%u/", heap_caps_get_largest_free_block(MALLOC_CAP_SPIRAM));
    MSG("%u\n", heap_caps_get_free_size(MALLOC_CAP_SPIRAM));
#endif
}

static void *spz_calloc(void *opaque, unsigned int items, unsigned int size)
{
    spz_stream *spz = opaque;
    heap_info();
    MSG("spz_calloc(%u,%u) = %u = ", items, size, items*size);
    void *p = calloc(items, size);
    CMSG("%p\n", p);
    heap_info();
    if (!p)
	spz->err = Z_MEM_ERROR;
    return p;
}
static void *spz_malloc(void *opaque, unsigned int size)
{
    spz_stream *spz = opaque;
    heap_info();
    MSG("spz_malloc(%u) = ", size);
    void *p = malloc(size);
    CMSG("%p\n", p);
    heap_info();
    if (!p)
	spz->err = Z_MEM_ERROR;
    return p;
}

static void spz_free(void *opaque, void *ptr)
{
    heap_info();
    MSG("spz_free(%p)\n", ptr);
    (void)opaque;
    free(ptr);
    heap_info();
}

int spz_read_data(spz_stream *spz, void *buf, size_t len)
{
    uint8_t *p = buf;

    while (len) {
	unsigned int avail = spz->zs.next_out - spz->optr;

	if (spz->err)
	    break;

	if (avail) {
	    if (avail > len)
		avail = len;

	    memcpy(p, spz->optr, avail);
	    p += avail;
	    spz->optr += avail;
	    len -= avail;
	} else {
	    spz->optr = spz->zs.next_out = spz->obuf;
	    spz->zs.avail_out = BUFFER_SIZE;

	    while (spz->zs.avail_out) {
		if (!spz->zs.avail_in && !spz->eoi) {
		    int rlen;

		    spz->zs.next_in = spz->ibuf;
		    rlen = spz->read_data(spz->token, spz->ibuf, BUFFER_SIZE);
		    if (rlen < 0) {
			if (!spz->err)
			    spz->err = rlen;
			rlen = 0;
		    }
		    spz->eoi = !rlen;
		    spz->zs.avail_in = rlen;
		}

		int rv = inflate(&spz->zs, Z_SYNC_FLUSH);
		if (rv == Z_OK || (rv == Z_BUF_ERROR && !spz->eoi))
		    continue;

		spz->eoi = true;
		if (rv != Z_STREAM_END && !spz->err)
		    spz->err = rv;
		break;
	    }
	}
    }
    return p - (uint8_t *)buf;
}

/*
 * spz needs to be initialized to zero except the read_data and cookie
 * fields.
 */
static int fwupdate_data_init(spz_stream *spz)
{
    spz->zs.zalloc = spz_calloc;
    spz->zs.zfree  = spz_free;
    spz->zs.opaque = spz;	/* For error reporting */
    spz->err = Z_OK;

    /* This is necessary due to unziplib damage */
    spz->zs.state = spz_calloc(spz, 1, sizeof(struct inflate_state));
    if (!spz->zs.state)
	goto err;

    for (int i = 0; i < SPZ_NBUF; i++) {
	spz->bufs[i] = spz_malloc(spz, BUFFER_SIZE);
	if (!spz->bufs[i])
	    goto err;
    }

    /* gzip, max window size */
    int rv = inflateInit2(&spz->zs, 16 + 15);
    printf("[FWUP] fwupdate_data_init: inflateInit2 returned %d\n", rv);
    if (rv != Z_OK && rv != Z_STREAM_END) {
	spz->err = rv;
	goto err;
    }
    spz->cleanup = true;

err:
    return spz->err;
}

static int fwupdate_data_cleanup(spz_stream *spz)
{
    int err = 0;

    if (!spz)
	return 0;

    err = spz->err;

    if (spz->cleanup)
	inflateEnd(&spz->zs);

    /* Don't reload the FPGA on error; it wedges the JTAG bus */
    if (spz->fpga_updated && !err)
	fpga_reset();

    for (int i = 0; i < SPZ_NBUF; i++) {
	if (spz->bufs[i])
	    free(spz->bufs[i]);
    }
    if (spz->zs.state)
	free(spz->zs.state);

    return err;
}

/*
 * Blash a full chunk of data as a JTAG SHIFT_DR transaction
 */
int jtag_shift_spz(spz_stream *spz, enum jtag_io_flags flags)
{
    unsigned int data_left = spz->header.len;
    int err = 0;

    if (!data_left)
	return 0;

    while (data_left) {
	unsigned int bytes = data_left;
	int rv;

	if (bytes > BUFFER_SIZE)
	    bytes = BUFFER_SIZE;

	rv = spz_read_data(spz, spz->dbuf, bytes);
	if (rv < 1) {
	    err = Z_DATA_ERROR;
	    break;
	}

	data_left -= rv;
	jtag_io(rv << 3, data_left ? 0 : flags, spz->dbuf, NULL);
    }

    return err;
}

static void *fwupdate_read_chunk_str(spz_stream *spz)
{
    int rv;

    if (spz->header.len >= BUFFER_SIZE) {
	spz->err = Z_DATA_ERROR;
	return NULL;
    }

    rv = spz_read_data(spz, spz->dbuf, spz->header.len);
    if (spz->err) {
	return NULL;
    }
    if (rv != (int)spz->header.len) {
	spz->err = Z_DATA_ERROR;
	return NULL;
    }
    spz->dbuf[spz->header.len] = '\0';
    return spz->dbuf;
}

/* Skip a data chunk */
static int fwupdate_skip_chunk(spz_stream *spz)
{
    unsigned int skip = spz->header.len;

    while (skip) {
	unsigned int block = skip;
	if (block > BUFFER_SIZE)
	    block = BUFFER_SIZE;

	int rv = spz_read_data(spz, spz->dbuf, block);
	if (spz->err)
	    return spz->err;
	if (rv != (int)block) {
	    return spz->err = Z_DATA_ERROR;
	}
	skip -= block;
    }

    return 0;
}

static int fwupdate_boardinfo(spz_stream *spz)
{
    uint8_t *board_info_data = spz_malloc(spz, BOARDINFO_SIZE);
    int rv = Z_OK;

    MSG("updating FPGA board_info\n");

    if (!board_info_data)
	return spz->err;

    if (board_info.len >= 16 &&
	board_info.len <= sizeof board_info &&
	board_info.len <= BOARDINFO_SIZE) {
	memcpy(board_info_data, &board_info, board_info.len);
	memset(board_info_data + board_info.len, 0xff,
	       BOARDINFO_SIZE - board_info.len);

	rv = spiflash_write_spz(spz, board_info_data, BOARDINFO_SIZE);
    }

    free(board_info_data);
    fwupdate_skip_chunk(spz);
    if (rv)
	spz->err = rv;
    return spz->err;
}

/* Process a data chunk; return a nonzero value if done */
static int fwupdate_process_chunk(spz_stream *spz)
{
    int rv;
    char *str;

    rv = spz_read_data(spz, &spz->header, sizeof spz->header);
    if (spz->err)
	return spz->err;
    else if (!rv)
	return Z_STREAM_END;
    else if (rv != sizeof spz->header)
	return spz->err = Z_STREAM_ERROR;

    if (spz->header.magic != FW_MAGIC) {
	MSG("bad chunk header magic 0x%08x\n", spz->header.magic);
	return spz->err = Z_DATA_ERROR;
    }

    switch (spz->header.type) {
    case FDT_END:
	return Z_STREAM_END;	/* End of data - not an error */
    case FDT_DATA:
	MSG("updating FPGA flash\n");
	return spiflash_write_spz(spz, NULL, 0);
    case FDT_BOARDINFO:
	return fwupdate_boardinfo(spz);
    case FDT_TARGET:
	str = fwupdate_read_chunk_str(spz);
#if 0
	if (!str || strcmp(str, spz->flash->target)) {
	    MSG("this firmware file targets \"%s\", need \"%s\"\n",
		str, spz->flash->target);
	    return spz->err = Z_DATA_ERROR;
	}
#else
	MSG("firmware target: \"%s\"\n", str);
#endif
	return Z_OK;
    case FDT_NOTE:
	str = fwupdate_read_chunk_str(spz);
	MSG("%s\n", str);
	return Z_OK;
    case FDT_ESP_OTA:
	MSG("updating ESP32... ");
	spz->esp_updated = true;
	rv = esp_update((read_func_t)spz_read_data, (token_t)spz,
			spz->header.len);
	CMSG("done.\n");
	return rv;
    case FDT_FPGA_INIT:
	MSG("initializing FPGA for flash programming... ");
	spz->fpga_updated = true;
	rv = fpga_program_spz(spz);
	CMSG("done\n");
	return rv;
    case FDT_ESP_PART:
    case FDT_ESP_SYS:
    case FDT_ESP_TOOL:
	/* Not applicable to this update method */
	return fwupdate_skip_chunk(spz);
    default:
	if (spz->header.flags & FDF_OPTIONAL) {
	    return fwupdate_skip_chunk(spz);
	} else {
	    MSG("unknown chunk type: %u\n", spz->header.type);
	    return spz->err = Z_DATA_ERROR;
	}
    }
}

const char *firmware_errstr(int err)
{
    static char unknown_err[32];
    static const char * const errstr[] = {
	[-Z_STREAM_ERROR]	         = "Decompression error",
	[-Z_DATA_ERROR]		         = "Invalid data stream",
	[-Z_MEM_ERROR]		         = "Out of memory",
	[-Z_BUF_ERROR]		         = "Decompression error",
	[-FWUPDATE_ERR_IN_PROGRESS]      =
	    "Firmware update already in progress",
	[-FWUPDATE_ERR_BAD_CHUNK]        = "Invalid firmware chunk header",
	[-FWUPDATE_ERR_ERASE_FAILED]     = "FPGA flash erase failed",
	[-FWUPDATE_ERR_PROGRAM_FAILED]   = "FGPA flash program failed",
	[-FWUPDATE_ERR_WRITE_PROTECT]    = "FPGA flash write protected",
	[-FWUPDATE_ERR_NOT_READY]        = "FPGA flash stuck at not ready",
	[-FWUPDATE_ERR_FPGA_JTAG]        =
	    "FPGA JTAG bus stuck, check for JTAG adapter or power cycle board",
	[-FWUPDATE_ERR_FPGA_MISMATCH]    =
	    "Bad FPGA IDCODE, check for JTAG adapter or power cycle board",
	[-FWUPDATE_ERR_FPGA_FAILED]      = "FPGA reboot failed",
	[-FWUPDATE_ERR_UNKNOWN]          = "Unidentified error",
	[-FWUPDATE_ERR_ESP_NO_PARTITION] = "No available ESP partition",
	[-FWUPDATE_ERR_ESP_BAD_OTA]      = "ESP OTA information corrupt",
	[-FWUPDATE_ERR_ESP_FLASH_FAILED] = "ESP flash program failed",
	[-FWUPDATE_ERR_ESP_BAD_DATA]     = "ESP firmware image corrupt",
	[-FWUPDATE_ERR_CONFIG_READ]      = "Configuration upload failure",
	[-FWUPDATE_ERR_CONFIG_SAVE]      = "Error saving configuration"
    };

    switch (err) {
    case Z_OK:
	return errstr[-FWUPDATE_ERR_UNKNOWN];
    case Z_ERRNO:
	return strerror(errno);
    case -ARRAY_SIZE(errstr)+1 ... Z_STREAM_ERROR:
	if (errstr[-err])
	    return errstr[-err];
	/* fall through */
    default:
	snprintf(unknown_err, sizeof unknown_err, "error %d", -err);
	return unknown_err;
    }
}

static TaskHandle_t fwupdate_task;
static spz_stream *fwupdate_spz;
static SemaphoreHandle_t fwupdate_done;
static int fwupdate_err;
static bool do_reboot;

static void firmware_update_task(void *pvt)
{
    spz_stream *spz = pvt;

    fpga_service_enable(false);

    printf("[FWUP] fwupdate_data_init()\n");

    spz->err = fwupdate_data_init(spz);
    if (spz->err)
	goto fail;

    printf("[FWUP] fwupdate_process_chunk loop\n");

    int err;
    while (!(err = fwupdate_process_chunk(spz))) {
	/* Process data chunks until end */
    }

    if (!spz->err && err != Z_STREAM_END)
	spz->err = err;

    printf("[FWUP] fwupdate_data_cleanup\n");

    err = fwupdate_data_cleanup(spz);
    if (err)
	spz->err = err;

fail:
    if (spz->err)
	MSG("failed (err %d)\n", spz->err);
    xSemaphoreGive(fwupdate_done);

    if (do_reboot) {
	printf("[FWUP] rebooting in %d seconds\n", reboot_delayed());
	while (1)
	    vTaskSuspend(NULL);
    } else {
	exit_task();
    }
}

static int firmware_update_cleanup(void)
{
    int err = Z_OK;

    fwupdate_task = NULL;

    if (fwupdate_done) {
	SemaphoreHandle_t done = fwupdate_done;
	fwupdate_done = NULL;
	vSemaphoreDelete(done);
    } else {
	err = Z_MEM_ERROR;
    }
    if (fwupdate_spz) {
	struct spz_stream *spz = fwupdate_spz;
	if (spz->err)
	    err = spz->err;
	fwupdate_spz = NULL;
	free(spz);
    } else {
	err = Z_MEM_ERROR;
    }

    return err;
}

int firmware_update_start(read_func_t read_data, token_t token, bool autoreboot)
{
    int err;
    SemaphoreHandle_t done = NULL;

    do_reboot = autoreboot;

    if (fwupdate_spz)
	return FWUPDATE_ERR_IN_PROGRESS;

    fwupdate_spz = calloc(1, sizeof *fwupdate_spz);
    if (!fwupdate_spz)
	goto err;

    fwupdate_spz->read_data = read_data;
    fwupdate_spz->token     = token;

    fwupdate_done = xSemaphoreCreateBinary();
    if (!fwupdate_done)
	goto err;

    if (xTaskCreate(firmware_update_task, "fwupdate",
		    FWUPDATE_STACK, fwupdate_spz,
		    FWUPDATE_PRIORITY, &fwupdate_task) != pdPASS) {
	xSemaphoreGive(fwupdate_done);
    }
    return Z_OK;

err:
    return firmware_update_cleanup();
}

int firmware_update_wait(TickType_t delay)
{
    if (!fwupdate_done)
	return Z_MEM_ERROR;

    if (!xSemaphoreTake(fwupdate_done, delay))
	return FWUPDATE_ERR_IN_PROGRESS;

    return firmware_update_cleanup();
}