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

#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 spiflash_header header;  /* Header of currently processed chunk */
    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) {
	    MSG("\rupdate: erasing %2uK at 0x%06x... ", erasesize >> 10, addr);
	
	    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;
	}

	/* 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", rv);
	    rv = Z_DATA_ERROR;
	    goto err;
	}

	MSG("update: flash block at 0x%06x (%5u bytes):\n", addr, bytes);

	addr -= pre_padding;

	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) {
	    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)
		goto err;
	    
	    spiflash_check_block(spz, addr, &erase_mask, prog_mask);
	    if (erase_mask) {
		MSG("[erase mask = %04x] ", erase_mask);
		spz->err = SPIFLASH_ERR_ERASE_FAILED;
		goto err;
	    }
	}

	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("\rupdate: 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 @ 0x%06x ", addr + page_offs);
		spz->err = SPIFLASH_ERR_PROGRAM_FAILED;
		goto err;
	    }
	}
	if (programmed)
	    MSG("ok\n");
	else
	    MSG("update: nothing to write\n");
	
	addr += pre_padding + bytes;
	data_left -= bytes;
    }
    return spz->err;

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

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

    rv = spiflash_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 != SPIFLASH_MAGIC) {
	MSG("update: 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:
	if (!spz->flash->ops)
	    return spiflash_skip_chunk(spz);
	else
	    return spiflash_flash_chunk(spz);
    case FDT_TARGET:
	str = spiflash_read_chunk_str(spz);
	if (!str || strcmp(str, spz->flash->target)) {
	    MSG("update: this firmware file targets \"%s\", need \"%s\"\n",
		str, spz->flash->target);
	    return spz->err = Z_DATA_ERROR;
	}
	return Z_OK;
    case FDT_NOTE:
	str = spiflash_read_chunk_str(spz);
	MSG("update: %s\n", str);
	return Z_OK;
    default:
	if (spz->header.flags & FDF_OPTIONAL) {
	    return spiflash_skip_chunk(spz);
	} else {
	    MSG("update: unknown chunk type: %u\n", spz->header.type);
	    return spz->err = Z_DATA_ERROR;
	}
    }
}

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;

    if (0 && flash->ops) {
	static const uint8_t read_sr_cmd[2] = { ROM_READ_SR1, ROM_READ_SR2 };
	uint8_t sr1;
	uint32_t *sfdp;

	sfdp = malloc(SPIFLASH_SFDP_SIZE);
	memset(sfdp, 0, SPIFLASH_SFDP_SIZE);

	/* Note: SFDP data is littleendian! */
	err = spiflash_get_sfdp(flash, sfdp);
	if (err)
	    return err;

	for (int i = 0; i < SPIFLASH_SFDP_SIZE; i += 16) {
	    MSG("%04x :", i);
	    for (int j = 0; j < 16; j += 4) {
		MSG(" %08x", sfdp[(i+j) >> 2]);
	    }
	    MSG("\n");
	}
    
	if (sfdp[0] != 0x50444653) {
	    MSG("update: invalid SFDP information read\n");
	    return SPIFLASH_ERR_DETECT;
	}

	/*
	 * If the flash is busy, try to reset it
	 */
	err = spiflash_get_status(flash, ROM_READ_SR1, &sr1);
	if (err)
	    return err;
	if (sr1 & 0x01) {
	    udelay(60);
	    err = spiflash_get_status(flash, ROM_READ_SR1, &sr1);
	    if (err)
		return err;
	    if (sr1 & 0x01) {
		MSG("update: flash busy, trying reset... ");
		
		err = spiflash_simple_command(flash, ROM_ENABLE_RESET);
		if (err)
		    return err;
		err = spiflash_simple_command(flash, ROM_RESET);
		if (err)
		    return err;
		
		udelay(60);
		err = spiflash_get_status(flash, ROM_READ_SR1, &sr1);
		if (err || (sr1 & 0x01)) {
		    MSG("failed\n");
		    return SPIFLASH_ERR_NOT_READY;
		}
		MSG("ok\n");
	    }
	}
    }

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