#include <stdlib.h>
#include <string.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "driver/rmt_tx.h"
#include "esp_log.h"
#include "led_strip.h"

static const char *TAG = "LED_STRIP";

// WS2812 Timing (in nanoseconds)
#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

// LED strip structure
struct led_strip_t {
    rmt_channel_handle_t channel;
    rmt_encoder_handle_t encoder;
    uint16_t led_count;
    uint8_t *pixel_buf;
};

// RMT encoder for WS2812
typedef struct {
    rmt_encoder_t base;
    rmt_encoder_t *bytes_encoder;
    rmt_encoder_t *copy_encoder;
    int state;
    rmt_symbol_word_t reset_code;
} ws2812_encoder_t;

static size_t ws2812_encode(rmt_encoder_t *encoder, rmt_channel_handle_t channel,
                           const void *primary_data, size_t data_size, rmt_encode_state_t *ret_state)
{
    ws2812_encoder_t *ws2812_encoder = __containerof(encoder, ws2812_encoder_t, base);
    rmt_encode_state_t session_state = RMT_ENCODING_RESET;
    rmt_encode_state_t state = RMT_ENCODING_RESET;
    size_t encoded_symbols = 0;

    switch (ws2812_encoder->state) {
        case 0: // Send RGB data
            encoded_symbols += ws2812_encoder->bytes_encoder->encode(ws2812_encoder->bytes_encoder, channel,
                                                                    primary_data, data_size, &session_state);
            if (session_state & RMT_ENCODING_COMPLETE) {
                ws2812_encoder->state = 1; // Switch to reset code
            }
            if (session_state & RMT_ENCODING_MEM_FULL) {
                state |= RMT_ENCODING_MEM_FULL;
                goto out;
            }
            // Fall through

        case 1: // Send reset code
            encoded_symbols += ws2812_encoder->copy_encoder->encode(ws2812_encoder->copy_encoder, channel,
                                                                   &ws2812_encoder->reset_code,
                                                                   sizeof(ws2812_encoder->reset_code), &session_state);
            if (session_state & RMT_ENCODING_COMPLETE) {
                ws2812_encoder->state = 0; // Back to RGB data
                state |= RMT_ENCODING_COMPLETE;
            }
            if (session_state & RMT_ENCODING_MEM_FULL) {
                state |= RMT_ENCODING_MEM_FULL;
                goto out;
            }
    }

out:
    *ret_state = state;
    return encoded_symbols;
}

static esp_err_t ws2812_del(rmt_encoder_t *encoder)
{
    ws2812_encoder_t *ws2812_encoder = __containerof(encoder, ws2812_encoder_t, base);
    rmt_del_encoder(ws2812_encoder->bytes_encoder);
    rmt_del_encoder(ws2812_encoder->copy_encoder);
    free(ws2812_encoder);
    return ESP_OK;
}

static esp_err_t ws2812_reset(rmt_encoder_t *encoder)
{
    ws2812_encoder_t *ws2812_encoder = __containerof(encoder, ws2812_encoder_t, base);
    rmt_encoder_reset(ws2812_encoder->bytes_encoder);
    rmt_encoder_reset(ws2812_encoder->copy_encoder);
    ws2812_encoder->state = 0;
    return ESP_OK;
}

static esp_err_t ws2812_encoder_new(rmt_encoder_handle_t *ret_encoder)
{
    ws2812_encoder_t *ws2812_encoder = calloc(1, sizeof(ws2812_encoder_t));
    if (!ws2812_encoder) {
        ESP_LOGE(TAG, "No memory for WS2812 encoder");
        return ESP_ERR_NO_MEM;
    }

    ws2812_encoder->base.encode = ws2812_encode;
    ws2812_encoder->base.del = ws2812_del;
    ws2812_encoder->base.reset = ws2812_reset;

    // Create byte encoder
    rmt_bytes_encoder_config_t bytes_encoder_config = {
        .bit0 = {
            .level0 = 1,
            .duration0 = WS2812_T0H_NS / 25,  // 25ns resolution
            .level1 = 0,
            .duration1 = WS2812_T0L_NS / 25,
        },
        .bit1 = {
            .level0 = 1,
            .duration0 = WS2812_T1H_NS / 25,
            .level1 = 0,
            .duration1 = WS2812_T1L_NS / 25,
        },
        .flags.msb_first = 1,
    };
    ESP_ERROR_CHECK(rmt_new_bytes_encoder(&bytes_encoder_config, &ws2812_encoder->bytes_encoder));

    // Create copy encoder for reset code
    rmt_copy_encoder_config_t copy_encoder_config = {};
    ESP_ERROR_CHECK(rmt_new_copy_encoder(&copy_encoder_config, &ws2812_encoder->copy_encoder));

    // Setup reset code
    ws2812_encoder->reset_code = (rmt_symbol_word_t) {
        .level0 = 0,
        .duration0 = WS2812_RESET_US * 40,  // 25ns resolution, 40 = 1000/25
        .level1 = 0,
        .duration1 = 0,
    };

    *ret_encoder = &ws2812_encoder->base;
    return ESP_OK;
}

led_strip_t* led_strip_init(uint8_t gpio, uint16_t led_count)
{
    led_strip_t *strip = calloc(1, sizeof(led_strip_t));
    if (!strip) {
        ESP_LOGE(TAG, "No memory for LED strip");
        return NULL;
    }

    strip->led_count = led_count;
    strip->pixel_buf = calloc(led_count, 3);  // 3 bytes per LED (RGB)
    if (!strip->pixel_buf) {
        ESP_LOGE(TAG, "No memory for pixel buffer");
        free(strip);
        return NULL;
    }

    // Configure RMT TX channel
    rmt_tx_channel_config_t tx_config = {
        .gpio_num = gpio,
        .clk_src = RMT_CLK_SRC_DEFAULT,
        .resolution_hz = 40000000,  // 40MHz, 25ns resolution
        .mem_block_symbols = 64,
        .trans_queue_depth = 4,
    };
    
    if (rmt_new_tx_channel(&tx_config, &strip->channel) != ESP_OK) {
        ESP_LOGE(TAG, "Failed to create RMT TX channel");
        free(strip->pixel_buf);
        free(strip);
        return NULL;
    }

    // Create WS2812 encoder
    if (ws2812_encoder_new(&strip->encoder) != ESP_OK) {
        ESP_LOGE(TAG, "Failed to create WS2812 encoder");
        rmt_del_channel(strip->channel);
        free(strip->pixel_buf);
        free(strip);
        return NULL;
    }

    // Enable RMT channel
    rmt_enable(strip->channel);

    ESP_LOGI(TAG, "LED strip initialized on GPIO %d with %d LEDs", gpio, led_count);
    return strip;
}

esp_err_t led_strip_set_pixel(led_strip_t *strip, uint16_t index, uint8_t red, uint8_t green, uint8_t blue)
{
    if (!strip || index >= strip->led_count) {
        return ESP_ERR_INVALID_ARG;
    }

    // WS2812 expects GRB order
    strip->pixel_buf[index * 3 + 0] = green;
    strip->pixel_buf[index * 3 + 1] = red;
    strip->pixel_buf[index * 3 + 2] = blue;

    return ESP_OK;
}

esp_err_t led_strip_refresh(led_strip_t *strip)
{
    if (!strip) {
        return ESP_ERR_INVALID_ARG;
    }

    rmt_transmit_config_t tx_config = {
        .loop_count = 0,
    };

    return rmt_transmit(strip->channel, strip->encoder, strip->pixel_buf, 
                       strip->led_count * 3, &tx_config);
}

esp_err_t led_strip_clear(led_strip_t *strip)
{
    if (!strip) {
        return ESP_ERR_INVALID_ARG;
    }

    memset(strip->pixel_buf, 0, strip->led_count * 3);
    return led_strip_refresh(strip);
}

void led_strip_deinit(led_strip_t *strip)
{
    if (strip) {
        rmt_disable(strip->channel);
        rmt_del_channel(strip->channel);
        rmt_del_encoder(strip->encoder);
        free(strip->pixel_buf);
        free(strip);
    }
}