#ifdef ESP32

#include "enabled.h"

#if FASTLED_RMT5

#include "fl/unused.h"


/*
 * SPDX-FileCopyrightText: 2022-2024 Espressif Systems (Shanghai) CO LTD
 *
 * SPDX-License-Identifier: Apache-2.0
 */
#include <stdlib.h>
#include <string.h>
#include <sys/cdefs.h>
#include "esp_log.h"
#include "esp_check.h"
#include "driver/rmt_tx.h"
#include "led_strip.h"
#include "led_strip_interface.h"
#include "led_strip_rmt_encoder.h"

#define LED_STRIP_RMT_DEFAULT_RESOLUTION 10000000 // 10MHz resolution
#define LED_STRIP_RMT_DEFAULT_TRANS_QUEUE_SIZE 4

#ifndef LED_STRIP_RMT_DEFAULT_MEM_BLOCK_SYMBOLS_MULTIPLIER
#if CONFIG_IDF_TARGET_ESP32C3
// Fixes https://github.com/FastLED/FastLED/issues/1846
#define LED_STRIP_RMT_DEFAULT_MEM_BLOCK_SYMBOLS_MULTIPLIER 2
#else
#define LED_STRIP_RMT_DEFAULT_MEM_BLOCK_SYMBOLS_MULTIPLIER 1
#endif
#endif

// the memory size of each RMT channel, in words (4 bytes)
#ifndef LED_STRIP_RMT_DEFAULT_MEM_BLOCK_SYMBOLS
#if CONFIG_IDF_TARGET_ESP32 || CONFIG_IDF_TARGET_ESP32S2
#define LED_STRIP_RMT_DEFAULT_MEM_BLOCK_SYMBOLS (LED_STRIP_RMT_DEFAULT_MEM_BLOCK_SYMBOLS_MULTIPLIER * 64)
#else
#define LED_STRIP_RMT_DEFAULT_MEM_BLOCK_SYMBOLS (LED_STRIP_RMT_DEFAULT_MEM_BLOCK_SYMBOLS_MULTIPLIER * 48)
#endif
#endif


static const char *TAG = "led_strip_rmt";

typedef struct {
    led_strip_t base;
    rmt_channel_handle_t rmt_chan;
    rmt_encoder_handle_t strip_encoder;
    rmt_transmit_config_t tx_conf; /*!< Transmit configuration, stored in the object to prevent stack access in interrupts */
    uint32_t strip_len;
    uint8_t bytes_per_pixel;
    led_color_component_format_t component_fmt;
    uint8_t *pixel_buf;
    bool pixel_buf_allocated_internally; /*!< Flag to track if pixel_buf was allocated by this driver */
} led_strip_rmt_obj;

static esp_err_t led_strip_rmt_set_pixel(led_strip_t *strip, uint32_t index, uint32_t red, uint32_t green, uint32_t blue)
{
    led_strip_rmt_obj *rmt_strip = __containerof(strip, led_strip_rmt_obj, base);
    ESP_RETURN_ON_FALSE(index < rmt_strip->strip_len, ESP_ERR_INVALID_ARG, TAG, "index out of maximum number of LEDs");

    led_color_component_format_t component_fmt = rmt_strip->component_fmt;
    uint32_t start = index * rmt_strip->bytes_per_pixel;
    uint8_t *pixel_buf = rmt_strip->pixel_buf;

    pixel_buf[start + component_fmt.format.r_pos] = red & 0xFF;
    pixel_buf[start + component_fmt.format.g_pos] = green & 0xFF;
    pixel_buf[start + component_fmt.format.b_pos] = blue & 0xFF;
    if (component_fmt.format.num_components > 3) {
        pixel_buf[start + component_fmt.format.w_pos] = 0;
    }

    return ESP_OK;
}

static esp_err_t led_strip_rmt_set_pixel_rgbw(led_strip_t *strip, uint32_t index, uint32_t red, uint32_t green, uint32_t blue, uint32_t white)
{
    led_strip_rmt_obj *rmt_strip = __containerof(strip, led_strip_rmt_obj, base);
    led_color_component_format_t component_fmt = rmt_strip->component_fmt;
    ESP_RETURN_ON_FALSE(index < rmt_strip->strip_len, ESP_ERR_INVALID_ARG, TAG, "index out of maximum number of LEDs");
    ESP_RETURN_ON_FALSE(component_fmt.format.num_components == 4, ESP_ERR_INVALID_ARG, TAG, "led doesn't have 4 components");

    uint32_t start = index * rmt_strip->bytes_per_pixel;
    uint8_t *pixel_buf = rmt_strip->pixel_buf;

    pixel_buf[start + component_fmt.format.r_pos] = red & 0xFF;
    pixel_buf[start + component_fmt.format.g_pos] = green & 0xFF;
    pixel_buf[start + component_fmt.format.b_pos] = blue & 0xFF;
    pixel_buf[start + component_fmt.format.w_pos] = white & 0xFF;

    return ESP_OK;
}

static esp_err_t led_strip_rmt_refresh_async(led_strip_t *strip)
{
    led_strip_rmt_obj *rmt_strip = __containerof(strip, led_strip_rmt_obj, base);

    ESP_RETURN_ON_ERROR(rmt_enable(rmt_strip->rmt_chan), TAG, "enable RMT channel failed");
    ESP_RETURN_ON_ERROR(rmt_transmit(rmt_strip->rmt_chan, rmt_strip->strip_encoder, rmt_strip->pixel_buf,
                                     rmt_strip->strip_len * rmt_strip->bytes_per_pixel, &rmt_strip->tx_conf), TAG, "transmit pixels by RMT failed");
    return ESP_OK;
}

static esp_err_t led_strip_rmt_wait_for_done(led_strip_t *strip)
{
    led_strip_rmt_obj *rmt_strip = __containerof(strip, led_strip_rmt_obj, base);
    ESP_RETURN_ON_ERROR(rmt_tx_wait_all_done(rmt_strip->rmt_chan, -1), TAG, "wait for RMT done failed");
    ESP_RETURN_ON_ERROR(rmt_disable(rmt_strip->rmt_chan), TAG, "disable RMT channel failed");
    return ESP_OK;
}

static esp_err_t led_strip_rmt_refresh(led_strip_t *strip)
{
    ESP_RETURN_ON_ERROR(led_strip_rmt_refresh_async(strip), TAG, "refresh async failed");
    ESP_RETURN_ON_ERROR(led_strip_rmt_wait_for_done(strip), TAG, "wait for done failed");
    return ESP_OK;
}

static esp_err_t led_strip_rmt_clear(led_strip_t *strip)
{
    led_strip_rmt_obj *rmt_strip = __containerof(strip, led_strip_rmt_obj, base);
    // Write zero to turn off all leds
    memset(rmt_strip->pixel_buf, 0, rmt_strip->strip_len * rmt_strip->bytes_per_pixel);
    return led_strip_rmt_refresh(strip);
}

static esp_err_t led_strip_rmt_del(led_strip_t *strip)
{
    led_strip_rmt_obj *rmt_strip = __containerof(strip, led_strip_rmt_obj, base);
    ESP_RETURN_ON_ERROR(rmt_del_channel(rmt_strip->rmt_chan), TAG, "delete RMT channel failed");
    ESP_RETURN_ON_ERROR(rmt_del_encoder(rmt_strip->strip_encoder), TAG, "delete strip encoder failed");
    if (rmt_strip->pixel_buf_allocated_internally) {
      free(rmt_strip->pixel_buf);
    }
    free(rmt_strip);
    return ESP_OK;
}

esp_err_t led_strip_new_rmt_device(const led_strip_config_t *led_config, const led_strip_rmt_config_t *rmt_config, led_strip_handle_t *ret_strip)
{
    led_strip_rmt_obj *rmt_strip = NULL;
    esp_err_t ret = ESP_OK;
    FASTLED_UNUSED(ret);
    ESP_GOTO_ON_FALSE(led_config && rmt_config && ret_strip, ESP_ERR_INVALID_ARG, err, TAG, "invalid argument");
    led_color_component_format_t component_fmt = led_config->color_component_format;
    // If R/G/B order is not specified, set default GRB order as fallback
    if (component_fmt.format_id == 0) {
        component_fmt = LED_STRIP_COLOR_COMPONENT_FMT_GRB;
    }
    // check the validation of the color component format
    uint8_t mask = 0;
    if (component_fmt.format.num_components == 3) {
        mask = BIT(component_fmt.format.r_pos) | BIT(component_fmt.format.g_pos) | BIT(component_fmt.format.b_pos);
        // Check for invalid values
        ESP_RETURN_ON_FALSE(mask == 0x07, ESP_ERR_INVALID_ARG, TAG, "invalid order argument");
    } else if (component_fmt.format.num_components == 4) {
        mask = BIT(component_fmt.format.r_pos) | BIT(component_fmt.format.g_pos) | BIT(component_fmt.format.b_pos) | BIT(component_fmt.format.w_pos);
        // Check for invalid values
        ESP_RETURN_ON_FALSE(mask == 0x0F, ESP_ERR_INVALID_ARG, TAG, "invalid order argument");
    } else {
        ESP_RETURN_ON_FALSE(false, ESP_ERR_INVALID_ARG, TAG, "invalid number of color components: %d", component_fmt.format.num_components);
    }
    // TODO: we assume each color component is 8 bits, may need to support other configurations in the future, e.g. 10bits per color component?
    uint8_t bytes_per_pixel = component_fmt.format.num_components;
    rmt_strip = calloc(1, sizeof(led_strip_rmt_obj));
    ESP_GOTO_ON_FALSE(rmt_strip, ESP_ERR_NO_MEM, err, TAG, "no mem for rmt strip");
    
    // Check if external pixel buffer is provided
    if (led_config->external_pixel_buf != NULL) {
        // Use the external pixel buffer provided by caller
        rmt_strip->pixel_buf = led_config->external_pixel_buf;
        rmt_strip->pixel_buf_allocated_internally = false;
    } else {
        // Allocate our own pixel buffer
        rmt_strip->pixel_buf = calloc(led_config->max_leds * bytes_per_pixel, sizeof(uint8_t));
        ESP_GOTO_ON_FALSE(rmt_strip->pixel_buf, ESP_ERR_NO_MEM, err, TAG, "no mem for pixel buffer");
        rmt_strip->pixel_buf_allocated_internally = true;
    }
    uint32_t resolution = rmt_config->resolution_hz ? rmt_config->resolution_hz : LED_STRIP_RMT_DEFAULT_RESOLUTION;

    // for backward compatibility, if the user does not set the clk_src, use the default value
    rmt_clock_source_t clk_src = RMT_CLK_SRC_DEFAULT;
    if (rmt_config->clk_src) {
        clk_src = rmt_config->clk_src;
    }
    size_t mem_block_symbols = LED_STRIP_RMT_DEFAULT_MEM_BLOCK_SYMBOLS;
    // override the default value if the user sets it
    if (rmt_config->mem_block_symbols) {
        mem_block_symbols = rmt_config->mem_block_symbols;
    }
    rmt_tx_channel_config_t rmt_chan_config = {
        .clk_src = clk_src,
        .gpio_num = led_config->strip_gpio_num,
        .mem_block_symbols = mem_block_symbols,
        .resolution_hz = resolution,
        .trans_queue_depth = LED_STRIP_RMT_DEFAULT_TRANS_QUEUE_SIZE,
        .flags.with_dma = rmt_config->flags.with_dma,
        .flags.invert_out = led_config->flags.invert_out,
    };
    ESP_GOTO_ON_ERROR(rmt_new_tx_channel(&rmt_chan_config, &rmt_strip->rmt_chan), err, TAG, "create RMT TX channel failed");

    led_strip_encoder_config_t strip_encoder_conf = {
        .resolution = resolution,
        .led_model = led_config->led_model,
        .timings = led_config->timings,
    };
    ESP_GOTO_ON_ERROR(rmt_new_led_strip_encoder(&strip_encoder_conf, &rmt_strip->strip_encoder), err, TAG, "create LED strip encoder failed");

    rmt_strip->component_fmt = component_fmt;
    rmt_strip->bytes_per_pixel = bytes_per_pixel;
    rmt_strip->strip_len = led_config->max_leds;
    rmt_strip->tx_conf = (rmt_transmit_config_t){
        .loop_count = 0,
    };
    rmt_strip->base.set_pixel = led_strip_rmt_set_pixel;
    rmt_strip->base.set_pixel_rgbw = led_strip_rmt_set_pixel_rgbw;
    rmt_strip->base.refresh = led_strip_rmt_refresh;
    rmt_strip->base.refresh_async = led_strip_rmt_refresh_async;
    rmt_strip->base.refresh_wait_done = led_strip_rmt_wait_for_done;
    rmt_strip->base.clear = led_strip_rmt_clear;
    rmt_strip->base.del = led_strip_rmt_del;

    *ret_strip = &rmt_strip->base;
    return ESP_OK;
err:
    if (rmt_strip) {
        if (rmt_strip->rmt_chan) {
            rmt_del_channel(rmt_strip->rmt_chan);
        }
        if (rmt_strip->strip_encoder) {
            rmt_del_encoder(rmt_strip->strip_encoder);
        }
        if (rmt_strip->pixel_buf_allocated_internally) {
            free(rmt_strip->pixel_buf);
        }
        free(rmt_strip);
    }
    return ret;
}

#endif  // FASTLED_RMT5

#endif  // ESP32
