#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_READ_STATUS   0x05u
#define FLASH_CMD_WRITE_DISABLE 0x04u
#define FLASH_CMD_WRITE_ENABLE  0x06u
#define FLASH_CMD_READ_CONFIG   0x35u
#define FLASH_CMD_READ_4B       0x13u
#define FLASH_CMD_PROGRAM_4B    0x12u
#define FLASH_CMD_ERASE_4B      0xdcu

#define FLASH_STATUS_WIP (1u << 0)
#define FLASH_STATUS_WEL (1u << 1)

#if BSP_SPI_FLASH_PROFILE == BSP_SPI_FLASH_PROFILE_S25FL256S
#define FLASH_CMD_CLEAR_STATUS 0x30u
#define FLASH_CMD_RESET        0xf0u
#define FLASH_STATUS_ERROR     ((1u << 5) | (1u << 6))
#elif BSP_SPI_FLASH_PROFILE == BSP_SPI_FLASH_PROFILE_MT25QU01G
#define FLASH_CMD_RESET_ENABLE      0x66u
#define FLASH_CMD_RESET_MEMORY      0x99u
#define FLASH_CMD_READ_FLAG_STATUS  0x70u
#define FLASH_CMD_CLEAR_FLAG_STATUS 0x50u
#define FLASH_FLAG_STATUS_ERROR     ((1u << 1) | (1u << 4) | (1u << 5))
#elif BSP_SPI_FLASH_PROFILE == BSP_SPI_FLASH_PROFILE_IS25WP512M
#define FLASH_CMD_RESET_ENABLE 0x66u
#define FLASH_CMD_RESET_MEMORY 0x99u
#endif

#define FLASH_RESET_DELAY_US        300u
#define FLASH_PROGRAM_TIMEOUT_US 100000u
#define FLASH_ERASE_TIMEOUT_US  3000000u

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 void flash_delay_us(uint32_t microseconds)
{
  uint64_t ticks =
      ((uint64_t)BSP_CLINT_TIMEBASE_HZ * microseconds + 999999u) / 1000000u;
  uint64_t start = bsp_read64(BSP_CLINT_MTIME);
  while ((bsp_read64(BSP_CLINT_MTIME) - start) < ticks) {
  }
}

static void flash_send_address(uint32_t flash_addr)
{
  (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);
}

static uint8_t flash_read_register(uint8_t command)
{
  uint8_t value;
  spi_select(1);
  (void)spi_xfer8(command);
  value = spi_xfer8(0u);
  spi_select(0);
  return value;
}

static bool flash_range_valid(uint32_t flash_addr, uint32_t len)
{
  return (flash_addr < BSP_SPI_FLASH_SIZE) &&
         (len <= (BSP_SPI_FLASH_SIZE - flash_addr));
}

static bool flash_wait_complete(uint32_t timeout_us)
{
  uint64_t timeout_ticks =
      ((uint64_t)BSP_CLINT_TIMEBASE_HZ * timeout_us + 999999u) / 1000000u;
  uint64_t start = bsp_read64(BSP_CLINT_MTIME);
  uint8_t status;

  do {
    status = bsp_spi_flash_read_status();
#if BSP_SPI_FLASH_PROFILE == BSP_SPI_FLASH_PROFILE_S25FL256S
    if ((status & FLASH_STATUS_ERROR) != 0u) {
      /* S25FL-S deliberately holds WIP high after a program/erase error
       * until CLSR is issued.  Clear the latched failure so a later request
       * can recover, while still reporting this operation as failed. */
      flash_cmd(FLASH_CMD_CLEAR_STATUS);
      return false;
    }
#endif
    if ((status & FLASH_STATUS_WIP) == 0u) {
#if BSP_SPI_FLASH_PROFILE == BSP_SPI_FLASH_PROFILE_MT25QU01G
      uint8_t flags = flash_read_register(FLASH_CMD_READ_FLAG_STATUS);
      if ((flags & FLASH_FLAG_STATUS_ERROR) != 0u) {
        flash_cmd(FLASH_CMD_CLEAR_FLAG_STATUS);
        return false;
      }
#endif
      return true;
    }
  } while ((bsp_read64(BSP_CLINT_MTIME) - start) < timeout_ticks);

  return false;
}

static bool flash_write_enable(void)
{
#if BSP_SPI_FLASH_PROFILE == BSP_SPI_FLASH_PROFILE_S25FL256S
  flash_cmd(FLASH_CMD_CLEAR_STATUS);
#elif BSP_SPI_FLASH_PROFILE == BSP_SPI_FLASH_PROFILE_MT25QU01G
  flash_cmd(FLASH_CMD_CLEAR_FLAG_STATUS);
#endif
  flash_cmd(FLASH_CMD_WRITE_ENABLE);
#if BSP_SPI_FLASH_PROFILE == BSP_SPI_FLASH_PROFILE_S25FL256S
  return (bsp_spi_flash_read_status() &
          (FLASH_STATUS_WIP | FLASH_STATUS_WEL | FLASH_STATUS_ERROR)) ==
         FLASH_STATUS_WEL;
#else
  return (bsp_spi_flash_read_status() &
          (FLASH_STATUS_WIP | FLASH_STATUS_WEL)) == FLASH_STATUS_WEL;
#endif
}

static bool flash_verify(uint32_t flash_addr, const uint8_t *expected,
                         uint32_t len)
{
  uint8_t readback[BSP_SPI_FLASH_PAGE_SIZE];

  while (len != 0u) {
    uint32_t chunk = len > sizeof(readback) ? sizeof(readback) : len;
    bsp_spi_flash_read(flash_addr, readback, chunk);
    for (uint32_t i = 0u; i < chunk; ++i) {
      if (readback[i] != expected[i]) {
        return false;
      }
    }
    flash_addr += chunk;
    expected += chunk;
    len -= chunk;
  }
  return true;
}

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

void bsp_spi_flash_reset(void)
{
  spi_select(0);
#if BSP_SPI_FLASH_PROFILE == BSP_SPI_FLASH_PROFILE_S25FL256S
  flash_cmd(FLASH_CMD_RESET);
#else
  flash_cmd(FLASH_CMD_RESET_ENABLE);
  flash_cmd(FLASH_CMD_RESET_MEMORY);
#endif
  flash_delay_us(FLASH_RESET_DELAY_US);
}

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

uint8_t bsp_spi_flash_read_status(void)
{
  return flash_read_register(FLASH_CMD_READ_STATUS);
}

uint8_t bsp_spi_flash_read_config(void)
{
#if BSP_SPI_FLASH_PROFILE == BSP_SPI_FLASH_PROFILE_S25FL256S
  return flash_read_register(FLASH_CMD_READ_CONFIG);
#else
  return 0u;
#endif
}

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

  if ((len == 0u) || (dst == NULL) || !flash_range_valid(flash_addr, len)) {
    return;
  }

  spi_select(1);
  (void)spi_xfer8(FLASH_CMD_READ_4B);
  flash_send_address(flash_addr);
  for (uint32_t i = 0; i < len; ++i) {
    out[i] = spi_xfer8(0u);
  }
  spi_select(0);
}

bool bsp_spi_flash_erase_sector(uint32_t flash_addr)
{
  if ((flash_addr & (BSP_SPI_FLASH_SECTOR_SIZE - 1u)) != 0u ||
      !flash_range_valid(flash_addr, BSP_SPI_FLASH_SECTOR_SIZE)
#if BSP_SPI_FLASH_PROFILE == BSP_SPI_FLASH_PROFILE_S25FL256S
      ||
      ((flash_addr >= (BSP_SPI_FLASH_SIZE - (2u * BSP_SPI_FLASH_SECTOR_SIZE))) &&
       ((bsp_spi_flash_read_config() & BSP_SPI_FLASH_CONFIG_TBPARM) != 0u))
#endif
      ) {
    return false;
  }
  if (!flash_wait_complete(FLASH_PROGRAM_TIMEOUT_US) || !flash_write_enable()) {
    return false;
  }

  spi_select(1);
  (void)spi_xfer8(FLASH_CMD_ERASE_4B);
  flash_send_address(flash_addr);
  spi_select(0);

  if (!flash_wait_complete(FLASH_ERASE_TIMEOUT_US)) {
    flash_cmd(FLASH_CMD_WRITE_DISABLE);
    return false;
  }

  uint8_t erased[BSP_SPI_FLASH_PAGE_SIZE];
  for (uint32_t i = 0u; i < sizeof(erased); ++i) {
    erased[i] = 0xffu;
  }
  return flash_verify(flash_addr, erased, sizeof(erased)) &&
         flash_verify(flash_addr + BSP_SPI_FLASH_SECTOR_SIZE - sizeof(erased),
                      erased, sizeof(erased));
}

bool bsp_spi_flash_program(uint32_t flash_addr, const void *src, uint32_t len)
{
  const uint8_t *bytes = (const uint8_t *)src;

  if (len == 0u) {
    return true;
  }
  if ((src == NULL) || !flash_range_valid(flash_addr, len)) {
    return false;
  }

  while (len != 0u) {
    uint32_t page_remaining =
        BSP_SPI_FLASH_PAGE_SIZE - (flash_addr & (BSP_SPI_FLASH_PAGE_SIZE - 1u));
    uint32_t chunk = len < page_remaining ? len : page_remaining;

    if (!flash_wait_complete(FLASH_PROGRAM_TIMEOUT_US) ||
        !flash_write_enable()) {
      return false;
    }

    spi_select(1);
    (void)spi_xfer8(FLASH_CMD_PROGRAM_4B);
    flash_send_address(flash_addr);
    for (uint32_t i = 0u; i < chunk; ++i) {
      (void)spi_xfer8(bytes[i]);
    }
    spi_select(0);

    if (!flash_wait_complete(FLASH_PROGRAM_TIMEOUT_US) ||
        !flash_verify(flash_addr, bytes, chunk)) {
      flash_cmd(FLASH_CMD_WRITE_DISABLE);
      return false;
    }

    flash_addr += chunk;
    bytes += chunk;
    len -= chunk;
  }

  return true;
}