#pragma once

#include "esp_err.h"
#include "soc/rmt_struct.h"
#include "hal/rmt_types.h"
#include "driver/gpio.h"
#include <stdlib.h>

#include <stdlib.h>
#include <string.h>
#include <sys/cdefs.h>
#include "esp_log.h"
#include "esp_attr.h"
#include "driver/rmt.h"

#define WS2812_T0H_NS (350)
#define WS2812_T0L_NS (1000)
#define WS2812_T1H_NS (1000)
#define WS2812_T1L_NS (350)
#define WS2812_RESET_US (280)

static uint32_t ws2812_t0h_ticks = 0;
static uint32_t ws2812_t1h_ticks = 0;
static uint32_t ws2812_t0l_ticks = 0;
static uint32_t ws2812_t1l_ticks = 0;

static void IRAM_ATTR ws2812_rmt_adapter(const void *src, rmt_item32_t *dest, size_t src_size,
                                         size_t wanted_num, size_t *translated_size, size_t *item_num)
{
    if (src == NULL || dest == NULL)
    {
        *translated_size = 0;
        *item_num = 0;
        return;
    }
    const rmt_item32_t bit0 = {{{ws2812_t0h_ticks, 1, ws2812_t0l_ticks, 0}}}; // Logical 0
    const rmt_item32_t bit1 = {{{ws2812_t1h_ticks, 1, ws2812_t1l_ticks, 0}}}; // Logical 1
    size_t size = 0;
    size_t num = 0;
    uint8_t *psrc = (uint8_t *)src;
    rmt_item32_t *pdest = dest;
    while (size < src_size && num < wanted_num)
    {
        for (int i = 0; i < 8; i++)
        {
            // MSB first
            if (*psrc & (1 << (7 - i)))
            {
                pdest->val = bit1.val;
            }
            else
            {
                pdest->val = bit0.val;
            }
            num++;
            pdest++;
        }
        size++;
        psrc++;
    }
    *translated_size = size;
    *item_num = num;
}

class ws2812_driver
{
public:
    ws2812_driver(gpio_num_t pin, rmt_channel_t rmt_channel, uint16_t countLeds) : _rmt_channel(rmt_channel),
                                                                                   _pin(pin),
                                                                                   _countLeds(countLeds)
    {
    }

    ~ws2812_driver()
    {
        free(buffer);
        ESP_ERROR_CHECK(rmt_driver_uninstall(_rmt_channel));
    }

    esp_err_t init()
    {
        // Write zero to turn off all leds
        buffer = (uint8_t *)calloc(1, _countLeds * 3);
        memset(buffer, 0, _countLeds * 3);

        rmt_config_t config = RMT_DEFAULT_CONFIG_TX(_pin, _rmt_channel);
        // set counter clock to 40MHz
        config.clk_div = 2;

        ESP_ERROR_CHECK(rmt_config(&config));
        ESP_ERROR_CHECK(rmt_driver_install(_rmt_channel, 0, 0));

        uint32_t counter_clk_hz = 0;
        if (rmt_get_counter_clock(_rmt_channel, &counter_clk_hz) != ESP_OK)
        {
            ESP_LOGE("RMT_LED", "get rmt counter clock failed");
            return ESP_ERR_INVALID_ARG;
        }

        // ns -> ticks
        float ratio = (float)counter_clk_hz / 1e9;
        ws2812_t0h_ticks = (uint32_t)(ratio * WS2812_T0H_NS);
        ws2812_t0l_ticks = (uint32_t)(ratio * WS2812_T0L_NS);
        ws2812_t1h_ticks = (uint32_t)(ratio * WS2812_T1H_NS);
        ws2812_t1l_ticks = (uint32_t)(ratio * WS2812_T1L_NS);

        // set ws2812 to rmt adapter
        rmt_translator_init(_rmt_channel, ws2812_rmt_adapter);

        return ESP_OK;
    }

    esp_err_t set_pixel(uint16_t index, uint8_t red, uint8_t green, uint8_t blue)
    {
        if (index >= _countLeds)
            return ESP_ERR_INVALID_ARG;
        uint16_t start = index * 3;
        // In thr order of GRB
        buffer[start + 0] = green & 0xFF;
        buffer[start + 1] = red & 0xFF;
        buffer[start + 2] = blue & 0xFF;
        return ESP_OK;
    }

    esp_err_t update()
    {
        if (rmt_write_sample(_rmt_channel, buffer, _countLeds * 3, true) != ESP_OK)
        {
            ESP_LOGE("RMT_LED", "transmit RMT samples failed");
            return ESP_FAIL;
        }
        return rmt_wait_tx_done(_rmt_channel, pdMS_TO_TICKS(10));
    }

private:
    const rmt_channel_t _rmt_channel;
    const gpio_num_t _pin;
    const uint16_t _countLeds;
    uint8_t *buffer;
};