#include "common.h"
#include "zlib.h"
#include "spiflash.h"
#include "esp.h"
#include "matchver.h"
#include "boardinfo_fpga.h"

#if 1
# include "console.h"
# define MSG(...) con_printf(__VA_ARGS__)
#else
# define MSG(...) ((void)0)
#endif

struct spz_stream;
typedef struct spz_stream spz_stream;

#define NBUF 4
struct spz_stream {
    z_stream zs;
    const struct spiflash *flash;
    uint8_t *optr;		    /* Output data pointer into obuf */
    /* Note: available output data ends at zs->next_out */
    union {
	uint8_t *bufs[NBUF];
	struct {
	    uint8_t *ibuf;	    /* Input buffer if compressed */
	    uint8_t *obuf;	    /* Output buffer */
	    uint8_t *dbuf;	    /* Block data buffer */
	    uint8_t *vbuf;	    /* Readback/verify buffer */
	};
    };
    struct fw_header header;	    /* Header of currently processed chunk */
    struct fw_header vmatch;	    /* Matched firmware version string */
    int err;			    /* Error code to return */
    bool eoi;			    /* Reached end of input */
    bool cleanup;		    /* Call inflateEnd() */
};

static void *spz_malloc(spz_stream *spz, size_t bytes)
{
    void *p = malloc(bytes);
    if (!p && !spz->err) {
	spz->err = Z_MEM_ERROR;
    }
    return p;
}

static int spiflash_read_data(spz_stream *spz, void *buf, unsigned int 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 = SPIFLASH_BLOCK_SIZE;

	    while (spz->zs.avail_out) {
		if (!spz->zs.avail_in && !spz->eoi) {
		    int (*read_data)(void *, void *, unsigned int);
		    read_data = spz->flash->read_data;
		    spz->zs.next_in = spz->ibuf;
		    int rlen = read_data(spz->flash->cookie,
					 spz->ibuf, SPIFLASH_BLOCK_SIZE);

		    spz->eoi = rlen < SPIFLASH_BLOCK_SIZE;
		    if (rlen < 0) {
			if (!spz->err)
			    spz->err = rlen;
			rlen = 0;
		    }
		    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 flash, zs.next_in,
 * and zs.avail_in fields.
 */
static int spiflash_data_init(spz_stream *spz)
{
    int rv = Z_OK;
    uint8_t *rdbuf = NULL;
    int rlen;
    uint32_t header_crc;

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

    spz->eoi = !spz->flash->read_data;
    /* gzip, max window size */
    rv = inflateInit2(&spz->zs, 16 + 15);
    if (rv != Z_OK && rv != Z_STREAM_END) {
	spz->err = rv;
	goto err;
    }
    spz->cleanup = true;

err:
    return spz->err;
}

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

    if (!spz)
	return 0;

    err = spz->err;

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

    for (int i = 0; i < NBUF; i++) {
	if (spz->bufs[i])
	    free(spz->bufs[i]);
    }
    return err;
}

/*
 * Set up a command header with an address according to the SPI
 * addressing mode. Returns a pointer to the first byte past the
 * address.
 */
static void *spiflash_setup_addrcmd(const struct spiflash *flash,
				    uint32_t addr,
				    uint8_t cmd24, uint8_t cmd32,
				    void *cmdbuf)
{
    enum spiflash_addr_mode mode = flash->param->addr;
    uint8_t *cmd = cmdbuf;

    if (!mode)
	mode = addr < (1 << 24) ? SPIFLASH_ADDR_24BIT : SPIFLASH_ADDR_32BIT;

    if (mode == SPIFLASH_ADDR_24BIT) {
	*cmd++ = cmd24;
    } else {
	*cmd++ = cmd32;
	*cmd++ = addr >> 24;
    }
    *cmd++ = addr >> 16;
    *cmd++ = addr >> 8;
    *cmd++ = addr;

    return cmd;
}

static int spiflash_get_status(const struct spiflash *flash,
			       uint8_t cmd, uint8_t *sr)
{
    return flash->ops->spi_read(flash->cookie, &cmd, 1, sr, 1,
				flash->param->tshsl);
}

/* This needs a timeout function */
static int spiflash_wait_status(const struct spiflash *flash,
				int delay, uint8_t mask, uint8_t val)
{
    uint8_t sr1;
    int rv;

    do {
	if (flash->ops->yield)
	    flash->ops->yield(flash->cookie, delay);

	rv = spiflash_get_status(flash, ROM_READ_SR1, &sr1);
	if (rv)
	    return rv;
    } while ((sr1 & mask) != val); /* Waiting... */

    return 0;
}

int spiflash_read(const struct spiflash *flash,
		  uint32_t addr, void *buffer, size_t len)
{
    uint8_t cmdbuf[6];
    uint8_t *cmd;

    cmd = spiflash_setup_addrcmd(flash, addr,
				 ROM_FAST_READ, ROM_FAST_READ_32BIT,
				 cmdbuf);
    *cmd++ = 0;			/* Dummy cycles */

    return flash->ops->spi_read(flash->cookie, cmdbuf, cmd - cmdbuf,
				buffer, len, flash->param->tshsl1);
}

static int spiflash_simple_command(const struct spiflash *flash, uint8_t cmd)
{
    return flash->ops->spi_write(flash->cookie, &cmd, 1, NULL, 0,
				 flash->param->tshsl);
}

static int spiflash_write_enable(const struct spiflash *flash)
{
    uint8_t sr1;
    int rv;

    rv = spiflash_wait_status(flash, 0, 1, 0);
    if (rv)
	return rv;

    rv = spiflash_simple_command(flash, ROM_WRITE_ENABLE);
    if (rv)
	return rv;

    return spiflash_wait_status(flash, 0, 3, 2);
}

static int spiflash_program(const struct spiflash *flash,
			    uint32_t addr, const void *buffer, size_t len)
{
    uint8_t cmdbuf[5];
    uint8_t *cmd;
    int rv;

    rv = spiflash_write_enable(flash);
    if (rv)
	return rv;

    cmd = spiflash_setup_addrcmd(flash, addr,
				 ROM_PAGE_PROGRAM, ROM_PAGE_PROGRAM_32BIT,
				 cmdbuf);
    rv = flash->ops->spi_write(flash->cookie, cmdbuf, cmd - cmdbuf,
			       buffer, len, flash->param->tshsl2);
    if (rv)
	return rv;

    return spiflash_wait_status(flash, flash->param->tpp, 3, 0);
}

/*
 * Erase up to (long bits) sectors, using block erase if possible.
 */
static int spiflash_erase(const struct spiflash *flash,
			  uint32_t addr, unsigned long sector_mask)
{
    uint8_t cmdbuf[5];
    uint8_t *cmd;
    uint8_t cmd24, cmd32;
    uint32_t erasesize;
    int rv;
    int delay;
    const uint32_t block_mask = SPIFLASH_BLOCK_SIZE - 1;
    const unsigned long block_sector_mask
	= block_mask >> SPIFLASH_SECTOR_SHIFT;

    if (!sector_mask) {
	MSG("update: nothing to erase\n");
	return 0;
    }

    while (sector_mask) {
	if (!(addr & block_mask) &&
	    ((sector_mask & block_sector_mask) == block_sector_mask)) {
		cmd24 = ROM_ERASE_64K;
		cmd32 = ROM_ERASE_64K_32BIT;
		delay = flash->param->tbe2;
		erasesize = SPIFLASH_BLOCK_SIZE;
	} else {
	    cmd24 = ROM_ERASE_4K;
	    cmd32 = ROM_ERASE_4K_32BIT;
	    delay = flash->param->tse;
	    erasesize = SPIFLASH_SECTOR_SIZE;
	}

	if (sector_mask & 1) {
	    rv = spiflash_write_enable(flash);
	    if (rv)
		return rv;

	    cmd = spiflash_setup_addrcmd(flash, addr, cmd24, cmd32, cmdbuf);
	    rv = flash->ops->spi_write(flash->cookie, cmdbuf, cmd - cmdbuf,
				       NULL, 0, flash->param->tshsl2);
	    if (rv)
		return rv;

	    rv = spiflash_wait_status(flash, delay, 3, 0);
	    if (rv)
		return rv;
	}

	addr += erasesize;
	sector_mask >>= (erasesize >> SPIFLASH_SECTOR_SHIFT);
    }

    MSG("ok\n");
    return 0;
}

/*
 * from: current flash contents
 * to:   desired flash contents
 *
 * These are assumed to be aligned full block buffers
 */
enum flashmem_status {
    FMS_DONE,			/* All done, no programming needed */
    FMS_PROGRAM,		/* Can be programmed */
    FMS_ERASE,			/* Needs erase before programming */
    FMS_NOTCHECKED		/* Not checked yet */
};
static enum flashmem_status
spiflash_memcmp(const void *from, const void *to, size_t len)
{
    const uint32_t *pf = from;
    const uint32_t *pt = to;
    const uint32_t *pfend = (const uint32_t *)((const char *)from + len);
    uint32_t doprog  = 0;
    uint32_t doerase = 0;

    while (pf < pfend) {
	uint32_t f = *pf++;
	uint32_t t = *pt++;

	doprog  |=  f ^ t;	/* Need programming if any data mismatch */
	doerase |= ~f & t;	/* Need erasing if any 0 -> 1 */
    }

    return doerase ? FMS_ERASE : doprog ? FMS_PROGRAM : FMS_DONE;
}

/*
 * Check a block for sectors which need erasing and pages which need
 * programming; the prog_mask is 256 bits long and so span multiple words.
 *
 * The desired input is spz->dbuf and the existing flash content should be
 * already read into spz->vbuf.
 *
 */
static void spiflash_check_block(spz_stream *spz, uint32_t addr,
				uint32_t *erase_mask, uint32_t *prog_mask)
{
    const uint8_t *p, *q;
    unsigned int page;

    *erase_mask = 0;
    memset(prog_mask, 0, SPIFLASH_BLOCK_SIZE/SPIFLASH_PAGE_SIZE/8);

    p = spz->vbuf;
    q = spz->dbuf;

    for (page = 0; page < SPIFLASH_BLOCK_SIZE/SPIFLASH_PAGE_SIZE; page++) {
	enum flashmem_status status;

	switch (spiflash_memcmp(p, q, SPIFLASH_PAGE_SIZE)) {
	case FMS_ERASE:
	    *erase_mask |= UINT32_C(1) <<
		(page >> (SPIFLASH_SECTOR_SHIFT-SPIFLASH_PAGE_SHIFT));
	    break;
	case FMS_PROGRAM:
	    prog_mask[page >> 5] |= UINT32_C(1) << (page & 31);
	    break;
	default:
	    /* Nothing to do! */
	    break;
	}

	p += SPIFLASH_PAGE_SIZE;
	q += SPIFLASH_PAGE_SIZE;
    }
}

static int spiflash_flash_chunk(spz_stream *spz)
{
    unsigned int data_left = spz->header.len;
    unsigned int addr      = spz->header.addr;
    int rv;

    while (data_left && !spz->err) {
	unsigned int pre_padding  = addr & (SPIFLASH_BLOCK_SIZE-1);
	unsigned int post_padding;
	unsigned int bytes;

	bytes = SPIFLASH_BLOCK_SIZE - pre_padding;
	post_padding = 0;
	if (bytes > data_left) {
	    post_padding = bytes - data_left;
	    bytes = data_left;
	}
	addr -= pre_padding;

	/* Read the current content of this block into vbuf */
	rv = spiflash_read(spz->flash, addr, spz->vbuf, SPIFLASH_BLOCK_SIZE);
	if (rv)
	    goto err;

	/* Copy any invariant chunk */
	if (pre_padding)
	    memcpy(spz->dbuf, spz->vbuf, pre_padding);
	if (post_padding)
	    memcpy(spz->dbuf+SPIFLASH_BLOCK_SIZE-post_padding,
		   spz->vbuf+SPIFLASH_BLOCK_SIZE-post_padding,
		   post_padding);

	rv = spiflash_read_data(spz, spz->dbuf+pre_padding, bytes);
	if (rv != (int)bytes) {
	    MSG("needed %u bytes got %d, \n", bytes, rv);
	    rv = Z_DATA_ERROR;
	    goto err;
	}

	uint32_t erase_mask;
	uint32_t prog_mask[SPIFLASH_BLOCK_SIZE >> (SPIFLASH_PAGE_SHIFT+5)];

	spiflash_check_block(spz, addr, &erase_mask, prog_mask);
	if (erase_mask) {
	    MSG("flash: erasing at 0x%06x mask %04x... ", addr, erase_mask);

	    rv = spiflash_erase(spz->flash, addr, erase_mask);
	    if (rv)
		goto err;

	    /* Verify that the sector did erase */
	    rv = spiflash_read(spz->flash, addr, spz->vbuf, SPIFLASH_BLOCK_SIZE);
	    if (rv) {
		MSG("readback ");
		goto err;
	    }

	    spiflash_check_block(spz, addr, &erase_mask, prog_mask);
	    if (erase_mask) {
		spz->err = FWUPDATE_ERR_ERASE_FAILED;
		MSG("%04x left, ", erase_mask);
		goto err;
	    }
	    MSG("ok\n");
	}

	unsigned int page;
	bool programmed = false;

	for (page = 0; page < (SPIFLASH_BLOCK_SIZE >> SPIFLASH_PAGE_SHIFT);
	     page++) {
	    uint32_t page_offs = page << SPIFLASH_PAGE_SHIFT;

	    if (!(prog_mask[page >> 5] & (UINT32_C(1) << (page & 31))))
		continue;	/* No need to program */

	    programmed = true;

	    udelay(100);
	    MSG("flash: writing at 0x%06x... ", addr + page_offs);

	    rv = spiflash_program(spz->flash, addr + page_offs,
				  spz->dbuf + page_offs,
				  SPIFLASH_PAGE_SIZE);
	    if (rv)
		goto err;

	    /* Verify that the page did write */
	    rv = spiflash_read(spz->flash, addr + page_offs,
			       spz->vbuf + page_offs,
			       SPIFLASH_PAGE_SIZE);
	    if (rv) {
		MSG("readback ");
		goto err;
	    }

	    if (memcmp(spz->dbuf + page_offs, spz->vbuf + page_offs,
		       SPIFLASH_PAGE_SIZE)) {
		MSG("verify ");
		spz->err = FWUPDATE_ERR_PROGRAM_FAILED;
		goto err;
	    }
	    MSG("ok\n");
	}
	if (programmed)
	    MSG("ok\n");
	else
	    MSG("unchanged\n");

	addr += pre_padding + bytes;
	data_left -= bytes;
    }
    return spz->err;

err:
    MSG("failed\n");

    if (!spz->err)
	spz->err = rv;

    return spz->err;
}

/* Serial Flash Discoverable Parameter Table, see JESD216 */
static int spiflash_get_sfdp(const struct spiflash *flash, void *sfdp)
{
    static const uint8_t cmd_read_sfdp[] = { ROM_READ_SFDP, 0, 0, 0, 0 };

    return flash->ops->spi_read(flash->cookie, cmd_read_sfdp,
				sizeof cmd_read_sfdp, sfdp, SPIFLASH_SFDP_SIZE,
				flash->param->tshsl);
}

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

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

    rv = spiflash_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 spiflash_skip_chunk(spz_stream *spz)
{
    unsigned int skip = spz->header.len;

    while (skip) {
	unsigned int block = min(skip, SPIFLASH_BLOCK_SIZE);
	int rv = spiflash_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;
}

/* Read a chunk into malloc()'d storage */
static int spiflash_load_chunk(spz_stream *spz, void **dptr)
{
    void *data;
    int len = spz->header.len;

    *dptr = NULL;
    data = spz_malloc(spz, len);
    if (!data) {
	spiflash_skip_chunk(spz);
	return spz->err;
    } else {
	int rv = spiflash_read_data(spz, data, len);
	if (!spz->err && rv != len)
	    spz->err = Z_DATA_ERROR;
	if (spz->err) {
	    free(data);
	    return spz->err;
	}
	*dptr = data;
	return rv;
    }
}

static int esp_ota_chunk(spz_stream *spz)
{
    void *data;
    int len = spiflash_load_chunk(spz, &data);
    if (data) {
	esp_ota(data, len);
	free(data);
    }
    return spz->err;
}

/* Get a piece of the chunk header */
static int fwupdate_get_header_data(spz_stream *spz, void *buf, int len)
{
    int rv;

    rv = spiflash_read_data(spz, buf, len);
    if (spz->err)
	return spz->err;
    else if (!rv)
	return Z_STREAM_END;
    else if (rv != len)
	return spz->err = Z_STREAM_ERROR;
    else
	return Z_OK;
}

/* Get and validate a chunk header */
static int fwupdate_get_header(spz_stream *spz)
{
    struct fw_header * const hdr = &spz->header;
    uint8_t *hptr = (uint8_t *)hdr;
    int rv;
    unsigned int hlen;

    memset(hdr, 0, sizeof *hdr);
    hdr->vmax = -1;

    rv = fwupdate_get_header_data(spz, hptr, FW_HDR_LEN_V1);
    if (rv)
	return rv;

    switch (hdr->magic) {
    case FW_MAGIC_V1:
	hlen = FW_HDR_LEN_V1;
	break;

    case FW_MAGIC_V2:
	hlen = FW_HDR_LEN_V2;
	break;

    default:
	MSG("update: bad chunk header magic 0x%08x\n", hdr->magic);
	hlen = 0;
	rv = Z_DATA_ERROR;
	break;
    }

    if (hlen > FW_HDR_LEN_V1) {
	rv = fwupdate_get_header_data(spz, hptr + FW_HDR_LEN_V1,
				      hlen - FW_HDR_LEN_V1);

	if (rv == Z_STREAM_END)	/* Only valid for the first chunk */
	    rv = Z_STREAM_ERROR;
    }

    return spz->err = rv;
}

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

    rv = fwupdate_get_header(spz);
    if (rv)
	return rv;

    if (spz->header.type != FDT_NOTE &&
	spz->header.type != FDT_TARGET &&
	spz->header.type != FDT_END &&
	!(spz->header.flags & FDF_PRETARGET)) {
	if (!spz->vmatch.magic) {
	    /* No matching firmware target support */
	    return spz->err = FWUPDATE_ERR_NOT_MINE;
	}

	if (spz->header.vmin > spz->vmatch.vmax ||
	    spz->header.vmax < spz->vmatch.vmin ||
	    ((spz->header.vmatch ^ spz->vmatch.vmatch) & spz->header.vmask)) {
	    /* Chunk not applicable to this target */
	    goto skip;
	}
    }

    con_printf("update: chunk type %u size %u addr 0x%08x\n",
	       spz->header.type, spz->header.len, spz->header.addr);

    switch (spz->header.type) {
    case FDT_END:
	return Z_STREAM_END;	/* End of data - not an error */
    case FDT_DATA:
	if (!spz->flash->ops)
	    goto skip;
	return spiflash_flash_chunk(spz);
    case FDT_TARGET:
    {
	bool match;
	str = spiflash_read_chunk_str(spz);
	match = match_version(board_info.version_str, str);
	if (match || spz->header.magic == FW_MAGIC_V1)
	    spz->vmatch = spz->header;

	MSG("update: firmware file supports: %s%s\n",
	    str, match ? " (match)" : "");
	return Z_OK;
    }
    case FDT_NOTE:
	str = spiflash_read_chunk_str(spz);
	MSG("update: %s\n", str);
	break;
    case FDT_ESP_OTA:
	if (!spz->flash->ops)
	    goto skip;
	return esp_ota_chunk(spz);
    case FDT_FPGA_INIT:
    case FDT_ESP_PART:
    case FDT_ESP_SYS:
    case FDT_ESP_TOOL:
    case FDT_BOARDINFO:
	/* Used only when flashing from ESP32 */
	goto skip;
    default:
	if (spz->header.flags & FDF_OPTIONAL)
	    goto skip;

	MSG("update: unknown chunk type: %u\n", spz->header.type);
	return spz->err = Z_DATA_ERROR;
    }
    return spz->err;

skip:
    return spiflash_skip_chunk(spz);
}

int spiflash_flash_file(const struct spiflash *flash, void *buf, size_t buflen)
{
    spz_stream _spz;
    spz_stream * const spz = &_spz; /* For consistency in notation */
    int err = 0;

    memset(spz, 0, sizeof *spz);
    spz->zs.avail_in = buflen;
    spz->zs.next_in  = buf;
    spz->flash = flash;

    err = spiflash_data_init(spz);
    if (err)
	return err;

    while (!spiflash_process_chunk(spz)) {
	/* Process data chunks until end */
    }

    err = spiflash_data_cleanup(spz);
    if (err)
	MSG("failed (err %d)\n", err);

    return err;
}

/*
 * Read unique serial number from flash. Note: returns id in
 * bigendian ("network") byte order.
 */
int spiflash_read_id(const struct spiflash *flash, void *id)
{
    static const uint8_t read_unique_id[] = { ROM_READ_UNIQUE_ID, 0, 0, 0, 0 };

    return flash->ops->spi_read(flash->cookie, read_unique_id,
				sizeof read_unique_id,
				id, SPIFLASH_ID_LEN, flash->param->tshsl);
}

/*
 * Read vendor and device ID from flash.
 */
int spiflash_read_vdid(const struct spiflash *flash, void *vdid)
{
    static const uint8_t read_vdid[] = { ROM_MANUFACTURER_DEVICE_ID, 0, 0, 0 };

    return flash->ops->spi_read(flash->cookie, read_vdid,
				sizeof read_vdid,
				vdid, SPIFLASH_VDID_LEN, flash->param->tshsl);
}

/*
 * Write an absolute region to flash
 */
int spiflash_flash_data(const struct spiflash *flash, uint32_t addr,
			const void *data, size_t len)
{
    spz_stream _spz;
    spz_stream * const spz = &_spz; /* For consistency in notation */
    int err = 0;

    memset(spz, 0, sizeof *spz);

    /* No ibuf or obuf */
    for (int i = 2; i < NBUF; i++) {
	spz->bufs[i] = spz_malloc(spz, SPIFLASH_BLOCK_SIZE);
	if (!spz->bufs[i])
		goto err;
    }

    spz->flash       = flash;
    spz->optr        = (uint8_t *)data; /* OK to lose const here */
    spz->obuf        = (uint8_t *)data; /* OK to lose const here */
    spz->zs.next_out = (uint8_t *)data + len;
    spz->eoi         = true;
    spz->header.len  = len;
    spz->header.addr = addr;

    spiflash_flash_chunk(spz);

err:
    for (int i = 2; i < NBUF; i++)
	if (spz->bufs[i])
	    free(spz->bufs[i]);

    return spz->err;
}