#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_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,
	.read = spi_flash_read,
	.write = spi_flash_write,
	.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 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, uint8_t *jedec_id,
	size_t len)
{
	uint8_t cmd[1] = { SPI_FLASH_JEDEC_ID };
	uint8_t 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)
{
	uint8_t 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] = { SPI_FLASH_READ };
	struct spi_dev *spi_dev = dev->priv;

	spi_flash_addr(cmd, addr);

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

	return len;
}

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

	spi_flash_write_enable(dev);

	cmd[0] = SPI_FLASH_PAGE_PROGRAM;
	spi_flash_addr(cmd, addr);
	memcpy(cmd + 4, data, len);

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

	spi_flash_write_disable(dev);

	return len;
}

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