#include <drv/ws2812b.h>

struct WS2812B {
  spi_device_handle_t spi_handle;
};

ws2812b_t* ws2812b_init(spi_device_handle_t spi)
{
  ws2812b_t* ret = heap_caps_malloc(sizeof(ws2812b_t), MALLOC_CAP_DEFAULT);
  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)
{
  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_ = heap_caps_malloc(n * 9, MALLOC_CAP_DMA);
  ret->n_rgb = n;
  return ret;
}

esp_err_t ws2812b_write_sync(ws2812b_t* drv, ws2812b_buffer_t* buffer)
{
  esp_err_t err;

  compile(buffer);
  spi_transaction_t t = {0};
  spi_transaction_t* rt;

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

  err = spi_device_queue_trans(drv->spi_handle, &t, portMAX_DELAY);
  if (err != ESP_OK) {
    return err;
  }

  err = spi_device_get_trans_result(drv->spi_handle, &rt, portMAX_DELAY);
  if (err != ESP_OK) {
    return err;
  }

  return ESP_OK;
}