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

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

static int spi_flash_read(struct flash_dev *dev, uint32_t addr, void *data,
	size_t len);
static int 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, size_t len);

static struct flash_ops spi_flash_ops ={
	.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_set_cs_level(spi_dev, 0);
	spi_tx_rx(spi_dev, NULL, cmd, sizeof cmd);
	spi_set_cs_level(spi_dev, 1);
}

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_set_cs_level(spi_dev, 0);
	spi_tx_rx(spi_dev, NULL, cmd, sizeof cmd);
	spi_set_cs_level(spi_dev, 1);
}

static int 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_set_cs_level(spi_dev, 0);
	spi_tx_rx(spi_dev, NULL, cmd, sizeof cmd);
	spi_tx_rx(spi_dev, data, NULL, len);
	spi_set_cs_level(spi_dev, 1);

	return 0;
}

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

	spi_flash_write_enable(dev);
	spi_flash_addr(cmd, addr);

	spi_set_cs_level(spi_dev, 0);
	spi_tx_rx(spi_dev, NULL, cmd, sizeof cmd);
	spi_tx_rx(spi_dev, NULL, data, len);
	spi_set_cs_level(spi_dev, 1);

	spi_flash_write_disable(dev);

	return 0;
}

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

	spi_flash_write_enable(dev);
	addr = ROUND_UP(addr, 4 * KIB);

	for (; len; addr += size, len -= size) {
		if (IS_ROUND(addr, 64 * KIB) && IS_ROUND(len, 64 * KIB))
			size = 64 * KIB;
		else if (IS_ROUND(addr, 32 * KIB) && IS_ROUND(len, 32 * KIB))
			size = 32 * KIB;
		else if (IS_ROUND(addr, 4 * KIB) && IS_ROUND(len, 4 * KIB))
			size = 4 * KIB;
		else
			size = 0;

		switch (size) {
		case 64 * KIB: cmd[0] = SPI_FLASH_ERASE_64K; break;
		case 32 * KIB: cmd[0] = SPI_FLASH_ERASE_32K; break;
		case 4 * KIB: cmd[0] = SPI_FLASH_ERASE_4K; break;
		default: goto err_disable_write;
		};

		printf("addr: %lx; size: %u\n", addr, size);
		spi_flash_addr(cmd, addr);

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

err_disable_write:
	spi_flash_write_disable(dev);
	return 0;
}

static char *size_to_str(size_t size)
{
	char *s;
	int ret;

	if (size == 0) {
		ret = asprintf(&s, "0");
#if defined(TIB)
	} else if (size % TIB == 0) {
		ret = asprintf(&s, "%u TiB", size / TIB);
#endif
	} else if (size % GIB == 0) {
		ret = asprintf(&s, "%u GiB", size / GIB);
	} else if (size % MIB == 0) {
		ret = asprintf(&s, "%u MiB", size / MIB);
	} else if (size % KIB == 0) {
		ret = asprintf(&s, "%u KiB", size / KIB);
	} else {
		ret = asprintf(&s, "%u B", size);
	}

	if (ret < 0)
		return NULL;

	return s;
}

static int flash_detect(struct flash_dev *dev)
{
	char cmd[1] = { SPI_FLASH_JEDEC_ID };
	uint8_t data[4];
	struct spi_dev *spi_dev = dev->priv;
	char *s;
	size_t size;

	spi_set_cs_level(spi_dev, 0);
	spi_tx_rx(spi_dev, data, cmd, sizeof cmd);
	spi_tx_rx(spi_dev, data, NULL, sizeof data);
	spi_set_cs_level(spi_dev, 1);

	printf("JEDEC ID: %02x%02x%02x%02x\n",
		data[0], data[1], data[2], data[3]);

	size = SPI_FLASH_SIZE(data[2]);
	s = size_to_str(size);

	printf("Capacity: %u blocks (%s)\n",
		size / (4 * KIB), s);
	free(s);

	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;

	flash_detect(dev);

	return dev;

err_free_dev:
	free(dev);
	return NULL;
}

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

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