#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.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 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 ={
	.get_size = 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_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_get_jedec_id(struct flash_dev *dev, uint8_t *jedec_id,
	size_t len)
{
	uint8_t cmd[5] = { SPI_FLASH_JEDEC_ID, 0, 0, 0, 0 };
	uint8_t buf[5];
	struct spi_dev *spi_dev = dev->priv;

	spi_set_cs_level(spi_dev, 0);
	spi_tx_rx(spi_dev, buf, cmd, sizeof cmd);
	spi_set_cs_level(spi_dev, 1);

	memcpy(jedec_id, buf + 1, min(4, len));

	return 0;
}

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

	printf("%02x%02x%02x%02x\n", jedec_id[0], jedec_id[1], jedec_id[2], jedec_id[3]);

	return SPI_FLASH_SIZE(jedec_id[2]);
}

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

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;

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