#include "compiler.h"
#include "zlib.h"
#include "spiflash.h"

struct spz_stream;
typedef struct spz_stream spz_stream;
struct spz_stream {
    z_stream zs;
    const struct spiflash *flash;
    int (*read_data)(spz_stream *); /* Routine to get more data */
    int (*end_data)(z_stream *);    /* Termination routine for zlib */
    int err;			    /* Error code to return */
    bool eoi;			    /* Reached end of input */
    uint8_t *optr;		    /* Output data pointer into obuf */
    /* Note: available output data ends at zs->next_out */
    uint8_t *ibuf;		    /* Input buffer if compressed */
    uint8_t *obuf;		    /* Output buffer */
    uint8_t *vbuf;		    /* Readback/verify buffer */
};

static int spiflash_read_data(spz_stream *spz)
{
    int rv;

    if (spz->eoi)
	return 0;

    spz->zs.next_in  = NULL;
    spz->zs.avail_in = 0;

    rv = spz->flash->ops->read_data(spz->flash->cookie,
				    spz->ibuf, SPIFLASH_BLOCK_SIZE);
    if (rv != SPIFLASH_BLOCK_SIZE)
	spz->eoi = true;

    if (rv < 0) {
	if (!spz->err)
	    spz->err = rv;
	return 0;
    }

    spz->zs.next_in  = spz->ibuf;
    spz->zs.avail_in = rv;

    return rv;
}

static int read_data_raw(spz_stream *spz)
{
    int rlen;

    if (spz->eoi)
	return Z_STREAM_END;

    rlen = spiflash_read_data(spz);
    if (rlen < 0)
	return rlen;		/* Error! */

    spz->optr         = spz->ibuf;
    spz->zs.next_out  = spz->ibuf + rlen;
    spz->zs.avail_out = SPIFLASH_BLOCK_SIZE - rlen;

    return spz->eoi ? Z_STREAM_END : Z_OK;
}

static int read_data_inflate(spz_stream *spz)
{
    int rv = Z_STREAM_END;

    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 rlen = spiflash_read_data(spz);
	    spz->zs.next_in  = spz->ibuf;
	    spz->zs.avail_in = rlen;
	}

	rv = inflate(&spz->zs, Z_SYNC_FLUSH);

	if (rv == Z_OK || (rv == Z_BUF_ERROR && !spz->eoi))
	    continue;

	if (rv != Z_STREAM_END) {
	    if (!spz->err)
		spz->err = rv >= 0 ? Z_STREAM_ERROR : rv;
	}
	break;
    }

    if (rv != Z_OK) {
	/* Z_STREAM_END or error */
	inflateEnd(&spz->zs);
    }

    return rv;
}

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;
}

int spiflash_data_init(spz_stream *spz, const struct spiflash *flash)
{
    int rlen;
    int rv = Z_OK;
    uint8_t *rdbuf = NULL;

    memset(spz, 0, sizeof *spz);
    spz->flash = flash;

    spz->ibuf = spz_malloc(spz, SPIFLASH_BLOCK_SIZE);
    if (!spz->ibuf)
	goto err;

    spz->vbuf = spz_malloc(spz, SPIFLASH_BLOCK_SIZE);
    if (!spz->vbuf)
	goto err;

    rlen = spiflash_read_data(spz);

    if (rlen >= 14 && !memcmp(spz->ibuf, "\37\213\10", 3)) {
	/* It is a gzip file */
	spz->zs.next_in   = spz->ibuf;
	spz->zs.avail_in  = rlen;
	spz->read_data    = read_data_inflate;

	spz->obuf = spz_malloc(spz, SPIFLASH_BLOCK_SIZE);
	if (!spz->obuf)
	    goto err;

	rv = inflateInit2(&spz->zs, 16 + 15);	/* gzip, max window size */
	if (rv != Z_OK) {
	    spz->err = rv;
	    goto err;
	}
	spz->end_data = inflateEnd;
    } else {
	/* Assume it is a raw binary; input buffer is output buffer */
	spz->optr         = spz->ibuf;
	spz->zs.next_out  = spz->ibuf + rlen;
	spz->zs.avail_out = SPIFLASH_BLOCK_SIZE - rlen;
	rv = spz->eoi ? Z_STREAM_END : Z_OK;
    }

err:
    return spz->err ? spz->err : rv;
 }

static void spiflash_data_cleanup(spz_stream *spz)
{
    if (!spz)
	return;

    if (spz->end_data)
	spz->end_data(&spz->zs);
    if (spz->vbuf)
	free(spz->vbuf);
    if (spz->obuf)
	free(spz->obuf);
    if (spz->ibuf)
	free(spz->ibuf);
}

/*
 * 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) {
	addr <<= 8;
	*cmd++ = cmd24;
    } else {
	*cmd++ = cmd32;
	*cmd++ = addr >> 24;
    }
    *cmd++ = addr >> 16;
    *cmd++ = addr >> 8;
    *cmd++ = addr;

    return cmd;
}

static int spiflash_wait_ready(const struct spiflash *flash, int delay)
{
    uint8_t cmd = 0x05;		/* Read Status Register 1 */
    uint8_t sr1 = 0;
    int rv;

    do {
	flash->ops->yield(flash->cookie, delay);
	rv = flash->ops->spi_read(flash->cookie, &cmd, 1,
				  &sr1, 1, flash->param->tshsl);
	if (rv)
	    return rv;
    } while (sr1 & 0x01);	/* Busy bit set? */

    return 0;
}

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

    /*
     * 13h = Fast Read
     * 0Ch = Fast Read with 4-Byte Address
     */
    cmd = spiflash_setup_addrcmd(flash, addr, 0x13, 0x0c, cmdbuf);
    *cmd++ = 0;			/* Dummy bits */

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

static int spiflash_write_enable(const struct spiflash *flash)
{
    const uint8_t cmd = 0x06;	/* 06h = Write Enable */

    return flash->ops->spi_write(flash->cookie, &cmd, 1, NULL, 0,
				 flash->param->tshsl);
}

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;

    /*
     * 02h = Page Program
     * 12h = Page Program with 4-Byte Address
     */
    cmd = spiflash_setup_addrcmd(flash, addr, 0x02, 0x12, cmdbuf);
    rv = flash->ops->spi_write(flash->cookie, cmdbuf, cmd - cmdbuf,
			       buffer, len, flash->param->tshsl2);
    if (rv)
	return rv;

    return spiflash_wait_ready(flash, flash->param->tpp);
}

/*
 * 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 nextaddr;
    int rv;
    int delay;
    const uint32_t block_mask = SPIFLASH_BLOCK_SIZE - 1;
    const unsigned long block_sectors = block_mask >> SPIFLASH_SECTOR_SHIFT;

    while (sector_mask) {
	if (((sector_mask & block_sectors) == block_sectors) &&
	    !(addr & block_mask)) {
		/*
		 * D8h = 64KB Block Erase
		 * DCh = 64K Block Erase with 4-Byte Address
		 */
		cmd24 = 0xd8; cmd32 = 0xdc;
		delay = flash->param->tbe2;
		nextaddr = addr + SPIFLASH_BLOCK_SIZE;
		sector_mask >>= 16;
	} else {
	    if (sector_mask & 1) {
		/*
		 * 20h = Sector Erase
		 * 21h = Sector Erase with 4-Byte Address
		 */
		cmd24 = 0x20; cmd32 = 0x21;
		delay = flash->param->tse;
		nextaddr = addr + SPIFLASH_SECTOR_SIZE;
		sector_mask >>= 1;
	    } else {
		addr += SPIFLASH_SECTOR_SIZE;
		sector_mask >>= 1;
		continue;	/* Skip sector */
	    }
	}
	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_ready(flash, delay);
	if (rv)
	    return rv;

	addr = nextaddr;
    }

    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 notok     = 0;
    uint32_t notprog   = 0;

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

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

    return notprog ? FMS_ERASE : notok ? 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 input is spz->optr, and the existing flash content is written
 * to spz->vptr.
 *
 */
static int spiflash_check_block(spz_stream *spz, uint32_t addr,
				uint32_t *erase_mask, uint32_t *prog_mask)
{
    int rv;
    const uint8_t *p, *q;
    unsigned int page;

    rv = spiflash_read(spz->flash, addr, spz->vbuf, SPIFLASH_BLOCK_SIZE);
    if (rv) {
	if (!spz->err)
	    spz->err = rv;
	return rv;
    }

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

    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;
    }

    return 0;
}

int spiflash_flash_file(const struct spiflash *flash, uint32_t addr)
{
    spz_stream _spz;
    spz_stream * const spz = &_spz; /* For consistency in notation */
    int rv;
    bool eof;
    enum flashmem_status fs;

    rv = spiflash_data_init(spz, flash);
    eof = rv != Z_OK;

    while (!eof && !spz->err) {
	unsigned int bytes = spz->zs.next_out - spz->optr;
	unsigned int padding;

	if (bytes < SPIFLASH_BLOCK_SIZE) {
	    int rv;

	    rv = spz->read_data(spz);
	    eof = rv != Z_OK || spz->err;
	    bytes = spz->zs.next_out - spz->optr;
	    if (!bytes)
		break;

	    padding = -bytes & (SPIFLASH_BLOCK_SIZE-1);
	    if (padding > spz->zs.avail_out) {
		/* This should never happen */
		padding = spz->zs.avail_out;
	    }
	    if (padding) {
		memset(spz->zs.next_out, 0xff, padding);
		spz->zs.avail_out -= padding;
		bytes += padding;
	    }
	}

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

	rv = spiflash_check_block(spz, addr, &erase_mask, prog_mask);
	if (rv)
	    break;

	if (erase_mask) {
	    rv = spiflash_erase(spz->flash, addr, erase_mask);
	    if (rv)
		break;
	    rv = spiflash_check_block(spz, addr, &erase_mask, prog_mask);
	    if (spz->err)
		break;
	    if (erase_mask) {
		spz->err = SPIFLASH_ERR_ERASE_FAILED;
		break;
	    }
	}

	unsigned int page;

	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))) {
		rv = spiflash_program(spz->flash, addr + page_offs,
				      spz->optr + page_offs,
				      SPIFLASH_PAGE_SIZE);
		if (rv) {
		    spz->err = rv;
		    break;
		}

		/* Read back data and verify */
		rv = spiflash_read(spz->flash, addr+page_offs,
				   spz->vbuf + page_offs,
				   SPIFLASH_PAGE_SIZE);
		if (rv) {
		    spz->err = rv;
		    break;
		}

		if (memcmp(spz->optr + page_offs, spz->vbuf + page_offs,
			   SPIFLASH_PAGE_SIZE)) {
			spz->err = SPIFLASH_ERR_PROGRAM_FAILED;
			break;
		}
	    }
	}

	spz->optr += SPIFLASH_BLOCK_SIZE;
	addr      += SPIFLASH_BLOCK_SIZE;
    }

    if (spz->flash->ops->close_data)
	spz->flash->ops->close_data(spz->flash->cookie);

    spiflash_data_cleanup(spz);
    return spz->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[] = { 0x4b, 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[] = { 0x90, 0, 0, 0 };

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