#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <bitops.h>
#include <flash.h>
#include <macros.h>
#include <spi.h>
#include <spi_flash.h>

static size_t spi_flash_get_jedec_id(struct flash_dev *dev, char *jedec_id,
	size_t len);
static size_t spi_flash_get_size(struct flash_dev *dev);
static size_t spi_flash_read(struct flash_dev *dev, uint32_t addr, void *data,
	size_t len);
static size_t spi_flash_write(struct flash_dev *dev, uint32_t addr,
	const void *data, size_t len);
static int spi_flash_erase(struct flash_dev *dev, uint32_t addr);
static void spi_flash_release(struct flash_dev *dev);

static struct flash_ops spi_flash_ops = {
	.release = spi_flash_release,
	.get_size = spi_flash_get_size,
	.get_capacity = spi_flash_get_size,
	.get_jedec_id = spi_flash_get_jedec_id,
	.read = spi_flash_read,
	.write = spi_flash_write,
	.copy = default_flash_copy,
	.is_erased = default_flash_is_one,
	.erase = spi_flash_erase,
};

static void spi_flash_addr(char *cmd, uint32_t addr)
{
	cmd[1] = addr >> 16;
	cmd[2] = addr >> 8;
	cmd[3] = addr;
}

static char spi_flash_get_status1(struct flash_dev *dev)
{
	struct spi_dev *spi_dev = dev->priv;
	char cmd = SPI_FLASH_READ_SR1;
	char ret;

	spi_tx_rx(spi_dev, &ret, sizeof ret, &cmd, sizeof cmd);

	return ret;
}

static int spi_flash_wait(struct flash_dev *dev)
{
	while (spi_flash_get_status1(dev) & SPI_FLASH_SR1_BUSY);

	return 0;
}

static void spi_flash_write_enable(struct flash_dev *dev)
{
	char cmd[1] = { SPI_FLASH_WRITE_ENABLE };
	struct spi_dev *spi_dev = dev->priv;

	spi_tx_rx(spi_dev, NULL, 0, cmd, sizeof cmd);
}

static void spi_flash_write_disable(struct flash_dev *dev)
{
	char cmd[1] = { SPI_FLASH_WRITE_DISABLE };
	struct spi_dev *spi_dev = dev->priv;

	spi_tx_rx(spi_dev, NULL, 0, cmd, sizeof cmd);
}

static size_t spi_flash_get_jedec_id(struct flash_dev *dev, char *jedec_id,
	size_t len)
{
	char cmd[1] = { SPI_FLASH_JEDEC_ID };
	char buf[3];
	struct spi_dev *spi_dev = dev->priv;

	spi_tx_rx(spi_dev, buf, sizeof buf, cmd, sizeof cmd);

	memcpy(jedec_id, buf, min(sizeof buf, len));

	return len;
}

static size_t spi_flash_get_size(struct flash_dev *dev)
{
	char jedec_id[4];

	spi_flash_get_jedec_id(dev, jedec_id, sizeof jedec_id);

	return SPI_FLASH_SIZE(jedec_id[2]);
}

static size_t spi_flash_read(struct flash_dev *dev, uint32_t addr, void *data,
	size_t len)
{
	char cmd[4];
	char *buf = data;
	struct spi_dev *spi_dev = dev->priv;
	size_t nbytes, ret = 0;

	while (len) {
		cmd[0] = SPI_FLASH_READ;
		spi_flash_addr(cmd, addr);

		nbytes = min(len, 32);

		spi_tx_rx(spi_dev, buf, len, cmd, sizeof cmd);

		addr += nbytes;
		buf += nbytes;
		len -= nbytes;
		ret += nbytes;
	}

	return ret;
}

static size_t spi_flash_write(struct flash_dev *dev, uint32_t addr,
	const void *data, size_t len)
{
	char cmd[4 + 32];
	const char *buf = data;
	struct spi_dev *spi_dev = dev->priv;
	size_t nbytes, ret = 0;

	while (len) {
		spi_flash_write_enable(dev);

		cmd[0] = SPI_FLASH_PAGE_PROGRAM;
		spi_flash_addr(cmd, addr);

		nbytes = min(len, 32);
		nbytes = min(nbytes, align_up(addr + 1, ilog2(256)) - addr);

		memmove(cmd + 4, buf, nbytes);
		spi_tx_rx(spi_dev, NULL, 0, cmd, 4 + nbytes);
		spi_flash_wait(dev);

		addr += nbytes;
		buf += nbytes;
		len -= nbytes;
		ret += nbytes;
	}

	return ret;
}

static int spi_flash_erase(struct flash_dev *dev, uint32_t addr)
{
	char cmd[4];
	struct spi_dev *spi_dev = dev->priv;

	spi_flash_write_enable(dev);

	cmd[0] = SPI_FLASH_ERASE_4K;
	spi_flash_addr(cmd, addr << dev->log2_block_size);

	spi_tx_rx(spi_dev, NULL, 0, cmd, sizeof cmd);
	spi_flash_wait(dev);

	spi_flash_write_disable(dev);

	return 0;
}

struct flash_dev *flash_probe(void)
{
	struct flash_dev *dev;

	if (!(dev = malloc(sizeof *dev)))
		return NULL;

	if (!(dev->priv = spi_probe()))
		goto err_free_dev;

	dev->ops = &spi_flash_ops;
	dev->log2_block_size = ilog2(4 * KIB);

	return dev;

err_free_dev:
	free(dev);
	return NULL;
}

static void spi_flash_release(struct flash_dev *dev)
{
	if (!dev)
		return;

	spi_release(dev->priv);
	free(dev);
}