#include "bsp/spi_flash.h"

#include "bsp/mmio.h"
#include "bsp/soc.h"

#include <stddef.h>

#define SPI_CTRL  (BSP_SPI_BASE + 0x00u)
#define SPI_TX    (BSP_SPI_BASE + 0x04u)
#define SPI_RX    (BSP_SPI_BASE + 0x08u)
#define SPI_READY (BSP_SPI_BASE + 0x0cu)

#define SPI_CTRL_BUSY 0x1u
#define SPI_CTRL_CSN  0x2u

#define FLASH_CMD_WRITE_ENABLE 0x06u
#define FLASH_CMD_READ_STATUS  0x05u
#define FLASH_CMD_PAGE_PROGRAM_4B 0x12u
#define FLASH_CMD_SECTOR_ERASE_4B 0xdcu
#define FLASH_STATUS_WIP 0x01u
#define FLASH_STATUS_WEL 0x02u
#define FLASH_PAGE_SIZE 256u
#define FLASH_READY_POLL_LIMIT 5000000u

static void delay_cycles(uint32_t cycles)
{
  for (volatile uint32_t i = 0; i < cycles; ++i) {
    asm volatile ("" ::: "memory");
  }
}

static void spi_wait_ready(void)
{
  while ((bsp_read32(SPI_READY) & 1u) == 0u) {
  }
  while ((bsp_read32(SPI_CTRL) & SPI_CTRL_BUSY) != 0u) {
  }
}

static void spi_select(int selected)
{
  spi_wait_ready();
  bsp_write32(SPI_CTRL, selected ? 0u : SPI_CTRL_CSN);
}

static uint8_t spi_xfer8(uint8_t value)
{
  spi_wait_ready();
  bsp_write32(SPI_TX, ((uint32_t)value) << 24);
  spi_wait_ready();
  return (uint8_t)(bsp_read32(SPI_RX) & 0xffu);
}

static void flash_cmd(uint8_t cmd)
{
  spi_select(1);
  (void)spi_xfer8(cmd);
  spi_select(0);
}

static uint8_t flash_read_status(void)
{
  uint8_t status;
  spi_select(1);
  (void)spi_xfer8(FLASH_CMD_READ_STATUS);
  status = spi_xfer8(0u);
  spi_select(0);
  return status;
}

static bool flash_write_enable(void)
{
  flash_cmd(FLASH_CMD_WRITE_ENABLE);
  return (flash_read_status() & FLASH_STATUS_WEL) != 0u;
}

static bool flash_wait_operation(void)
{
  for(uint32_t poll = 0u; poll < FLASH_READY_POLL_LIMIT; ++poll) {
    if((flash_read_status() & FLASH_STATUS_WIP) == 0u) {
      return true;
    }
  }
  return false;
}

void bsp_spi_flash_init(void)
{
  bsp_write32(SPI_CTRL, SPI_CTRL_CSN);
}

void bsp_spi_flash_reset(void)
{
  spi_select(0);
  flash_cmd(0x66u);
  delay_cycles(1000);
  flash_cmd(0x99u);
  delay_cycles(200000);
}

uint32_t bsp_spi_flash_read_id(void)
{
  uint32_t id = 0;
  spi_select(1);
  (void)spi_xfer8(0x9fu);
  id |= ((uint32_t)spi_xfer8(0u)) << 16;
  id |= ((uint32_t)spi_xfer8(0u)) << 8;
  id |= (uint32_t)spi_xfer8(0u);
  spi_select(0);
  return id;
}

void bsp_spi_flash_read(uint32_t flash_addr, void *dst, uint32_t len)
{
  uint8_t *out = (uint8_t *)dst;

  spi_select(1);
  (void)spi_xfer8(0x13u);
  (void)spi_xfer8((uint8_t)(flash_addr >> 24));
  (void)spi_xfer8((uint8_t)(flash_addr >> 16));
  (void)spi_xfer8((uint8_t)(flash_addr >> 8));
  (void)spi_xfer8((uint8_t)flash_addr);
  for (uint32_t i = 0; i < len; ++i) {
    out[i] = spi_xfer8(0u);
  }
  spi_select(0);
}

bool bsp_spi_flash_program_page(
  uint32_t flash_addr, const void *src, uint32_t len)
{
  if((src == NULL) || (len == 0u) || (len > FLASH_PAGE_SIZE) ||
     ((flash_addr & (FLASH_PAGE_SIZE - 1u)) + len > FLASH_PAGE_SIZE) ||
     !flash_write_enable()) {
    return false;
  }

  const uint8_t *data = (const uint8_t *)src;
  spi_select(1);
  (void)spi_xfer8(FLASH_CMD_PAGE_PROGRAM_4B);
  (void)spi_xfer8((uint8_t)(flash_addr >> 24));
  (void)spi_xfer8((uint8_t)(flash_addr >> 16));
  (void)spi_xfer8((uint8_t)(flash_addr >> 8));
  (void)spi_xfer8((uint8_t)flash_addr);
  for(uint32_t i = 0u; i < len; ++i) {
    (void)spi_xfer8(data[i]);
  }
  spi_select(0);
  return flash_wait_operation();
}

bool bsp_spi_flash_erase_sector(uint32_t flash_addr)
{
  if(!flash_write_enable()) {
    return false;
  }

  spi_select(1);
  (void)spi_xfer8(FLASH_CMD_SECTOR_ERASE_4B);
  (void)spi_xfer8((uint8_t)(flash_addr >> 24));
  (void)spi_xfer8((uint8_t)(flash_addr >> 16));
  (void)spi_xfer8((uint8_t)(flash_addr >> 8));
  (void)spi_xfer8((uint8_t)flash_addr);
  spi_select(0);
  return flash_wait_operation();
}