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

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

/* Read commands */
#define CMD_READ_ARRAY_SLOW		0x03
#define CMD_READ_ID				0x9f

/* Write commands */
#define CMD_PAGE_PROGRAM		0x02
#define CMD_WRITE_DISABLE		0x04
#define CMD_WRITE_ENABLE		0x06

/* Erase commands */
#define CMD_ERASE_4K			0x20
#define CMD_ERASE_32K			0x52
#define CMD_ERASE_64K			0xD8
#define CMD_ERASE_CHIP			0x60

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] = { CMD_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] = { CMD_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] = { CMD_READ_ARRAY_SLOW };
	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] = { CMD_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] = CMD_ERASE_64K; break;
		case 32 * KIB: cmd[0] = CMD_ERASE_32K; break;
		case 4 * KIB: cmd[0] = CMD_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;
}