#include <drv/ws2812b.h>
#include <memory.h>

#define WS2812B_FLAG_DIRTY 1

struct WS2812B {
  spi_device_handle_t spi_handle;
  spi_transaction_t spi_transaction;

  uint8_t flags;
};

ws2812b_t* ws2812b_init(spi_device_handle_t spi)
{
  ws2812b_t* ret = heap_caps_malloc(sizeof(ws2812b_t), MALLOC_CAP_DEFAULT);
  memset(ret, 0, sizeof(*ret));
  ret->spi_handle = spi;

  return ret;
}

/* Compiles a byte into it's padded serial format. */
#undef BIT
#define BIT(b, n) (!!((b) & (1 << (n))))
static inline void byte_compile(uint8_t byte, uint8_t compl [3])
{
  compl [0] = 0 | 1 << 7 | BIT(byte, 7) << 6 | 0 << 5 | 1 << 4 |
              BIT(byte, 6) << 3 | 0 << 2 | 1 << 1 | BIT(byte, 5) << 0;
  compl [1] = 0 | 0 << 7 | 1 << 6 | BIT(byte, 4) << 5 | 0 << 4 | 1 << 3 |
              BIT(byte, 3) << 2 | 0 << 1 | 1 << 0;
  compl [2] = 0 | BIT(byte, 2) << 7 | 0 << 6 | 1 << 5 | BIT(byte, 1) << 4 |
              0 << 3 | 1 << 2 | BIT(byte, 0) << 1 | 0 << 0;
}

static inline void compile(ws2812b_buffer_t* buffer)
{
  if (buffer->buf_ == buffer->buf_1) {
    buffer->buf_ = buffer->buf_2;
  } else {
    buffer->buf_ = buffer->buf_1;
  }

  for (size_t i = 0; i < buffer->n_rgb; ++i) {
    byte_compile(buffer->rgb[i].g, buffer->buf_ + i * 9 + 0);
    byte_compile(buffer->rgb[i].r, buffer->buf_ + i * 9 + 3);
    byte_compile(buffer->rgb[i].b, buffer->buf_ + i * 9 + 6);
  }
}

ws2812b_buffer_t* ws2812b_new_buffer(uint32_t n)
{
  ws2812b_buffer_t* ret = heap_caps_malloc(
      sizeof(ws2812b_buffer_t) + sizeof(ws2812b_rgb_t) * n, MALLOC_CAP_DEFAULT);
  ret->buf_1 = heap_caps_malloc(n * 9, MALLOC_CAP_DMA);
  ret->buf_2 = heap_caps_malloc(n * 9, MALLOC_CAP_DMA);
  ret->buf_ = ret->buf_1;
  ret->n_rgb = n;
  return ret;
}

esp_err_t ws2812b_wait(ws2812b_t* drv)
{
  esp_err_t err;

  spi_transaction_t* rt;
  err = spi_device_get_trans_result(
      drv->spi_handle, &rt, 10 / portTICK_PERIOD_MS);
  drv->flags &= ~WS2812B_FLAG_DIRTY;
  if (err != ESP_OK) {
    return err;
  }

  return ESP_OK;
}

esp_err_t ws2812b_write(ws2812b_t* drv, ws2812b_buffer_t* buffer)
{
  compile(buffer);

  if (drv->flags & WS2812B_FLAG_DIRTY) {
    ws2812b_wait(drv);
  }

  esp_err_t err;

  memset(&drv->spi_transaction, 0, sizeof(drv->spi_transaction));

  drv->spi_transaction.tx_buffer = buffer->buf_;
  drv->spi_transaction.length = buffer->n_rgb * 9 * 8;

  err = spi_device_queue_trans(
      drv->spi_handle, &drv->spi_transaction, portMAX_DELAY);
  if (err != ESP_OK) {
    return err;
  }
  drv->flags |= WS2812B_FLAG_DIRTY;

  return ESP_OK;
}