#include "compiler.h"
#include "zlib.h"
#include "spiflash.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;
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 */
    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 */
    const struct spiflash_header *header;
    uint32_t crc32;		    /* Input data CRC32 */
    unsigned int input_left;	    /* Input data unread */
    int err;			    /* Error code to return */
    bool eoi;			    /* Reached end of input */
    bool free_header;		    /* header is malloc()'d */
};

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)
{
    int rv;
    int (*read_data)(void *, void *, unsigned int);
    unsigned int read_block_size;

    if (spz->eoi || spz->err)
	return 0;

    read_data = spz->flash->ops->read_data;
    if (!spz->input_left || !read_data) {
	spz->eoi = true;
	return 0;
    }

    read_block_size = min(spz->input_left, SPIFLASH_BLOCK_SIZE);

    if (!spz->ibuf) {
	spz->ibuf = spz_malloc(spz, SPIFLASH_BLOCK_SIZE);
	if (!spz->ibuf) {
	    spz->eoi = true;
	    return 0;
	}
    }

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

    rv = read_data(spz->flash->cookie, spz->ibuf, read_block_size);
    if (spz->err) {
	rv = 0;
    } else if (rv < 0) {
	spz->err = rv;
	rv = 0;
    }
    if (rv != (int)read_block_size)
	spz->eoi = true;

    if (rv) {
	spz->crc32 = crc32(spz->crc32, spz->ibuf, rv);
	spz->input_left -= rv;
	if (!spz->input_left) {
	    if (spz->crc32 != spz->header->crc32) {
		spz->err = Z_STREAM_ERROR;
		rv = 0;
	    }
	}
    }

    return spz->zs.avail_in = rv;
}

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

    if (spz->eoi)
	return 0;

    rlen = spiflash_read_data(spz);
    if (rlen) {
	spz->optr         = spz->ibuf;
	spz->zs.next_out  = spz->ibuf + rlen;
	spz->zs.avail_out = SPIFLASH_BLOCK_SIZE - rlen;
    }

    return rlen;
}

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;

	spz->eoi = true;
	if (rv != Z_STREAM_END && !spz->err)
	    spz->err = rv;
	break;
    }

    return spz->zs.next_out - spz->optr;
}

/*
 * spz needs to be initialized to zero except the flash, zs.next_in,
 * and zs.avail_in fields.
 *
 * Returns Z_STREAM_END on end of data.
 */
static int spiflash_data_init(spz_stream *spz)
{
    int rv = Z_OK;
    uint8_t *rdbuf = NULL;
    int rlen;
    int (*read_data)(void *cookie, void *buf, unsigned int bufsize);
    uint32_t header_crc;

    MSG("update: initializing... ");

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

    rlen = spz->zs.avail_in;
    spz->header = (void *)spz->zs.next_in;

    read_data = spz->flash->ops->read_data;
    spz->eoi = !read_data;
    if (!rlen && read_data) {
	struct spiflash_header *header;
	spz->header = header = spz_malloc(spz, sizeof *spz->header);
	spz->free_header = true;
	spz->input_left = UINT_MAX; /* Unknown at this point */
	rlen = read_data(spz->flash->cookie, header, sizeof *header);
    }
    if (spz->err)
	goto err;
    if (rlen < (int)sizeof *spz->header) {
	spz->err = rlen ? Z_STREAM_ERROR : Z_STREAM_END;
	goto err;
    }

    rlen -= sizeof *spz->header;
    spz->zs.next_in += sizeof *spz->header;

    /*
     * Check header magic and CRC
     */
    if (spz->header->magic != SPIFLASH_MAGIC) {
	spz->err = Z_STREAM_ERROR;
	goto err;
    }
    header_crc = crc32(0, NULL, 0);
    header_crc = crc32(header_crc, (const void *)spz->header,
		       sizeof spz->header - 4);
    if (header_crc != spz->header->header_crc32 ||
	spz->header->zlen > spz->header->dlen ||
	(spz->eoi && (int)spz->header->zlen > rlen)) {
	spz->err = Z_STREAM_ERROR;
	goto err;
    }
    if (!spz->header->dlen) {
	/* End of data */
	spz->err = Z_STREAM_END;
	goto err;
    }

    MSG("updating 0x%06x..%06x (%u bytes)... ",
	spz->header->address, spz->header->address + spz->header->dlen - 1,
	spz->header->dlen);

    if (rlen > (int)spz->header->zlen)
	rlen = spz->header->zlen;
    spz->zs.avail_in = rlen;

    spz->crc32 = crc32(0, NULL, 0);
    if (rlen) {
	/* Received data in input buffer already */
	spz->crc32 = crc32(spz->crc32, spz->zs.next_in, spz->zs.avail_in);
    }

    spz->input_left = spz->header->zlen - rlen;
    if (!spz->input_left) {
	if (spz->crc32 != spz->header->crc32) {
	    spz->err = Z_STREAM_ERROR;
	    goto err;
	}
    }

    if (spz->header->zlen == spz->header->dlen) {
	/* Assume it is a raw binary; input buffer is output buffer */
	spz->read_data    = read_data_raw;
	spz->optr         = spz->zs.next_in;
	spz->zs.next_out  = spz->zs.next_in + spz->zs.avail_in;
    } else {
	/* Compressed data? */
	spz->obuf = spz_malloc(spz, SPIFLASH_BLOCK_SIZE);
	if (!spz->obuf)
	    goto err;

	if (rlen >= 14 && !memcmp(spz->zs.next_in, "\37\213\10", 3)) {
	    /* It is a gzip file */
	    spz->read_data = read_data_inflate;
	    /* gzip, max window size */
	    rv = inflateInit2(&spz->zs, 16 + 15);
	    if (rv != Z_OK && rv != Z_STREAM_END) {
		spz->err = rv;
		goto err;
	    }
	    spz->eoi = rv == Z_STREAM_END;
	    spz->end_data = inflateEnd;
	} else {
	    /* Unknown compression format */
	    spz->err = Z_STREAM_ERROR;
	    goto err;
	}
    }

err:
    MSG("%s\n", spz->err ? "failed" : "ok");
    return spz->err;
 }

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

    if (!spz)
	return 0;

    err = spz->err;

    if (spz->flash->ops->close_data) {
	int rv = spz->flash->ops->close_data(spz->flash->cookie);
	if (!err)
	    err = rv;
    }

    if (spz->end_data)
	spz->end_data(&spz->zs);
    if (spz->free_header)
	free((void *)spz->header);
    if (spz->vbuf)
	free(spz->vbuf);
    if (spz->obuf)
	free(spz->obuf);
    if (spz->ibuf)
	free(spz->ibuf);

    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) {
	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_files(const struct spiflash *flash, void *buf, size_t buflen)
{
    spz_stream _spz;
    spz_stream * const spz = &_spz; /* For consistency in notation */
    int  err;
    enum flashmem_status fs;
    uint8_t *padbuf = NULL;

    err = 0;
    while (!err) {
	int rv;
	bool eof;
	uint32_t addr;

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

	if (spiflash_data_init(spz))
	    goto err;

	if (!spz->ibuf) {
	    /* No ibuf allocated, feeding from raw data buffer */
	    unsigned int bufskip = (spz->zs.next_out - (uint8_t *)buf + 3) & ~3;
	    if (bufskip >= buflen) {
		buflen = 0;
		buf = NULL;
	    } else {
		buflen -= bufskip;
		buf = spz->zs.next_out;
	    }
	} else {
	    /* Buffer exhausted, additional data read */
	    buflen = 0;
	    buf = NULL;
	}

	eof = false;
	addr = spz->header->address;

	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);
		if (spz->err)
		    goto err;
		eof = !rv;
		bytes = spz->zs.next_out - spz->optr;
		if (!bytes)
		    break;

		padding = -bytes & (SPIFLASH_BLOCK_SIZE-1);
		if (padding) {
		    if (!padbuf) {
			padbuf = spz_malloc(spz, SPIFLASH_BLOCK_SIZE);
			if (!padbuf)
			    goto err;
		    }
		    memcpy(padbuf, spz->optr, bytes);
		    memset(padbuf+bytes, 0xff, padding);
		    spz->optr = padbuf;
		    eof = true;
		}
	    }

	    MSG("update: flash block at 0x%06x: ", addr);

	    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)
		goto err;

	    if (erase_mask) {
		MSG("erasing... ");

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

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

	    unsigned int page;
	    bool first_prog = true;

	    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))) {

		    if (first_prog) {
			MSG("programming %06x... ", addr + page_offs);
		    } else {
			MSG("\b\b\b\b\b\b\b\b\b\b%06x... ", addr+page_offs);
		    }
		    first_prog = false;

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

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

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

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

    err:
	err = spiflash_data_cleanup(spz);
    }

    if (padbuf) {
	free(padbuf);
	padbuf = NULL;
    }

    MSG("%s\n", spz->err ? "failed" : "ok");

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

/*
 * Flash data from memory buffer(s)
 */