/**
 * @file led.c
 * @brief WS2812B LED strip control implementation using RMT
 */

#include "led.h"

#include "driver/rmt_tx.h"
#include "esp_log.h"
#include "freertos/FreeRTOS.h"
#include "freertos/semphr.h"

#include <string.h>
#include <stdlib.h>

static const char *TAG = "LED";

// WS2812B timing (in nanoseconds)
#define WS2812_T0H_NS 350
#define WS2812_T0L_NS 900
#define WS2812_T1H_NS 900
#define WS2812_T1L_NS 350
#define WS2812_RESET_US 280

// LED strip data structures
typedef struct
{
    rmt_channel_handle_t rmt_channel;
    rmt_encoder_handle_t encoder;
    rgb_t *buffer;
    uint16_t num_leds;
    int8_t gpio_pin;
    bool initialized;
} led_strip_t;

static led_strip_t strip_a = {0};
static led_strip_t strip_b = {0};
static SemaphoreHandle_t led_mutex = NULL;

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

static size_t rmt_encode_led_strip(rmt_encoder_t *encoder, rmt_channel_handle_t channel,
                                   const void *primary_data, size_t data_size,
                                   rmt_encode_state_t *ret_state)
{
    rmt_led_strip_encoder_t *led_encoder = __containerof(encoder, rmt_led_strip_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 (led_encoder->state)
    {
    case 0: // send RGB data
        encoded_symbols += led_encoder->bytes_encoder->encode(led_encoder->bytes_encoder, channel,
                                                              primary_data, data_size, &session_state);
        if (session_state & RMT_ENCODING_COMPLETE)
        {
            led_encoder->state = 1; // switch to next state when current encoding session finished
        }
        if (session_state & RMT_ENCODING_MEM_FULL)
        {
            state |= RMT_ENCODING_MEM_FULL;
            goto out;
        }
    // fall-through
    case 1: // send reset code
        encoded_symbols += led_encoder->copy_encoder->encode(led_encoder->copy_encoder, channel,
                                                             &led_encoder->reset_code,
                                                             sizeof(led_encoder->reset_code), &session_state);
        if (session_state & RMT_ENCODING_COMPLETE)
        {
            led_encoder->state = RMT_ENCODING_RESET;
            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 rmt_del_led_strip_encoder(rmt_encoder_t *encoder)
{
    rmt_led_strip_encoder_t *led_encoder = __containerof(encoder, rmt_led_strip_encoder_t, base);
    rmt_del_encoder(led_encoder->bytes_encoder);
    rmt_del_encoder(led_encoder->copy_encoder);
    free(led_encoder);
    return ESP_OK;
}

static esp_err_t rmt_led_strip_encoder_reset(rmt_encoder_t *encoder)
{
    rmt_led_strip_encoder_t *led_encoder = __containerof(encoder, rmt_led_strip_encoder_t, base);
    rmt_encoder_reset(led_encoder->bytes_encoder);
    rmt_encoder_reset(led_encoder->copy_encoder);
    led_encoder->state = RMT_ENCODING_RESET;
    return ESP_OK;
}

static esp_err_t rmt_new_led_strip_encoder(rmt_encoder_handle_t *ret_encoder)
{
    esp_err_t ret = ESP_OK;
    rmt_led_strip_encoder_t *led_encoder = calloc(1, sizeof(rmt_led_strip_encoder_t));
    if (!led_encoder)
    {
        return ESP_ERR_NO_MEM;
    }

    led_encoder->base.encode = rmt_encode_led_strip;
    led_encoder->base.del = rmt_del_led_strip_encoder;
    led_encoder->base.reset = rmt_led_strip_encoder_reset;

    // WS2812 timing
    rmt_bytes_encoder_config_t bytes_encoder_config = {
        .bit0 = {
            .level0 = 1,
            .duration0 = WS2812_T0H_NS * 80 / 1000, // 80MHz clock
            .level1 = 0,
            .duration1 = WS2812_T0L_NS * 80 / 1000,
        },
        .bit1 = {
            .level0 = 1,
            .duration0 = WS2812_T1H_NS * 80 / 1000,
            .level1 = 0,
            .duration1 = WS2812_T1L_NS * 80 / 1000,
        },
        .flags.msb_first = 1,
    };
    ret = rmt_new_bytes_encoder(&bytes_encoder_config, &led_encoder->bytes_encoder);
    if (ret != ESP_OK)
    {
        goto err;
    }

    rmt_copy_encoder_config_t copy_encoder_config = {};
    ret = rmt_new_copy_encoder(&copy_encoder_config, &led_encoder->copy_encoder);
    if (ret != ESP_OK)
    {
        goto err;
    }

    uint32_t reset_ticks = WS2812_RESET_US * 80; // 80MHz
    led_encoder->reset_code = (rmt_symbol_word_t){
        .level0 = 0,
        .duration0 = reset_ticks & 0x7FFF,
        .level1 = 0,
        .duration1 = reset_ticks & 0x7FFF,
    };

    *ret_encoder = &led_encoder->base;
    return ESP_OK;

err:
    if (led_encoder->bytes_encoder)
    {
        rmt_del_encoder(led_encoder->bytes_encoder);
    }
    if (led_encoder->copy_encoder)
    {
        rmt_del_encoder(led_encoder->copy_encoder);
    }
    free(led_encoder);
    return ret;
}

static esp_err_t init_strip(led_strip_t *strip, int8_t pin, uint16_t num_leds)
{
    if (pin < 0 || num_leds == 0)
    {
        return ESP_OK; // Skip if not configured
    }

    strip->buffer = calloc(num_leds, sizeof(rgb_t));
    if (!strip->buffer)
    {
        return ESP_ERR_NO_MEM;
    }

    strip->num_leds = num_leds;
    strip->gpio_pin = pin;

    rmt_tx_channel_config_t tx_chan_config = {
        .clk_src = RMT_CLK_SRC_DEFAULT,
        .gpio_num = pin,
        .mem_block_symbols = 48,
        .resolution_hz = 80000000, // 80MHz
        .trans_queue_depth = 4,
    };

    ESP_ERROR_CHECK(rmt_new_tx_channel(&tx_chan_config, &strip->rmt_channel));
    ESP_ERROR_CHECK(rmt_new_led_strip_encoder(&strip->encoder));
    ESP_ERROR_CHECK(rmt_enable(strip->rmt_channel));

    strip->initialized = true;
    ESP_LOGI(TAG, "Initialized strip on GPIO%d with %d LEDs", pin, num_leds);

    return ESP_OK;
}

esp_err_t led_init(int8_t pin_a, int8_t pin_b, uint16_t num_leds_a, uint16_t num_leds_b)
{
    if (led_mutex == NULL)
    {
        led_mutex = xSemaphoreCreateMutex();
        if (!led_mutex)
        {
            return ESP_ERR_NO_MEM;
        }
    }

    esp_err_t ret = ESP_OK;

    if (pin_a >= 0)
    {
        ret = init_strip(&strip_a, pin_a, num_leds_a);
        if (ret != ESP_OK)
        {
            ESP_LOGE(TAG, "Failed to init strip A: %s", esp_err_to_name(ret));
            return ret;
        }
    }

    if (pin_b >= 0)
    {
        ret = init_strip(&strip_b, pin_b, num_leds_b);
        if (ret != ESP_OK)
        {
            ESP_LOGE(TAG, "Failed to init strip B: %s", esp_err_to_name(ret));
            return ret;
        }
    }

    return ESP_OK;
}

void led_deinit(void)
{
    if (strip_a.initialized)
    {
        rmt_disable(strip_a.rmt_channel);
        rmt_del_channel(strip_a.rmt_channel);
        free(strip_a.buffer);
        strip_a.initialized = false;
    }

    if (strip_b.initialized)
    {
        rmt_disable(strip_b.rmt_channel);
        rmt_del_channel(strip_b.rmt_channel);
        free(strip_b.buffer);
        strip_b.initialized = false;
    }
}

void led_set_pixel_a(uint16_t index, rgb_t color)
{
    if (!strip_a.initialized || index >= strip_a.num_leds)
        return;
    xSemaphoreTake(led_mutex, portMAX_DELAY);
    strip_a.buffer[index] = color;
    xSemaphoreGive(led_mutex);
}

void led_set_pixel_b(uint16_t index, rgb_t color)
{
    if (!strip_b.initialized || index >= strip_b.num_leds)
        return;
    xSemaphoreTake(led_mutex, portMAX_DELAY);
    strip_b.buffer[index] = color;
    xSemaphoreGive(led_mutex);
}

void led_fill_a(rgb_t color)
{
    if (!strip_a.initialized)
        return;
    xSemaphoreTake(led_mutex, portMAX_DELAY);
    for (uint16_t i = 0; i < strip_a.num_leds; i++)
    {
        strip_a.buffer[i] = color;
    }
    xSemaphoreGive(led_mutex);
}

void led_fill_b(rgb_t color)
{
    if (!strip_b.initialized)
        return;
    xSemaphoreTake(led_mutex, portMAX_DELAY);
    for (uint16_t i = 0; i < strip_b.num_leds; i++)
    {
        strip_b.buffer[i] = color;
    }
    xSemaphoreGive(led_mutex);
}

void led_clear_all(void)
{
    rgb_t black = {0, 0, 0};
    led_fill_a(black);
    led_fill_b(black);
}

static void show_strip(led_strip_t *strip)
{
    if (!strip->initialized)
        return;

    // Convert RGB to GRB for WS2812B
    uint8_t *grb_data = malloc(strip->num_leds * 3);
    if (!grb_data)
    {
        ESP_LOGE(TAG, "Failed to allocate GRB buffer");
        return;
    }

    for (uint16_t i = 0; i < strip->num_leds; i++)
    {
        grb_data[i * 3 + 0] = strip->buffer[i].g;
        grb_data[i * 3 + 1] = strip->buffer[i].r;
        grb_data[i * 3 + 2] = strip->buffer[i].b;
    }

    rmt_transmit_config_t tx_config = {
        .loop_count = 0,
    };

    esp_err_t ret = rmt_transmit(strip->rmt_channel, strip->encoder, grb_data, strip->num_leds * 3, &tx_config);
    if (ret != ESP_OK)
    {
        ESP_LOGE(TAG, "RMT transmit failed: %s", esp_err_to_name(ret));
        free(grb_data);
        return;
    }

    // Wait for transmission to complete before freeing buffer
    ret = rmt_tx_wait_all_done(strip->rmt_channel, pdMS_TO_TICKS(100));
    if (ret != ESP_OK)
    {
        ESP_LOGW(TAG, "RMT wait timeout");
    }

    free(grb_data);
}

void led_show(void)
{
    xSemaphoreTake(led_mutex, portMAX_DELAY);
    show_strip(&strip_a);
    show_strip(&strip_b);
    xSemaphoreGive(led_mutex);
}

void led_fade_to_black(uint8_t amount)
{
    xSemaphoreTake(led_mutex, portMAX_DELAY);

    if (strip_a.initialized)
    {
        for (uint16_t i = 0; i < strip_a.num_leds; i++)
        {
            strip_a.buffer[i].r = (strip_a.buffer[i].r * (255 - amount)) / 255;
            strip_a.buffer[i].g = (strip_a.buffer[i].g * (255 - amount)) / 255;
            strip_a.buffer[i].b = (strip_a.buffer[i].b * (255 - amount)) / 255;
        }
    }

    if (strip_b.initialized)
    {
        for (uint16_t i = 0; i < strip_b.num_leds; i++)
        {
            strip_b.buffer[i].r = (strip_b.buffer[i].r * (255 - amount)) / 255;
            strip_b.buffer[i].g = (strip_b.buffer[i].g * (255 - amount)) / 255;
            strip_b.buffer[i].b = (strip_b.buffer[i].b * (255 - amount)) / 255;
        }
    }

    xSemaphoreGive(led_mutex);
}

rgb_t led_hsv_to_rgb(hsv_t hsv)
{
    rgb_t rgb = {0};
    uint8_t region, remainder, p, q, t;

    if (hsv.s == 0)
    {
        rgb.r = hsv.v;
        rgb.g = hsv.v;
        rgb.b = hsv.v;
        return rgb;
    }

    region = hsv.h / 43;
    remainder = (hsv.h - (region * 43)) * 6;

    p = (hsv.v * (255 - hsv.s)) >> 8;
    q = (hsv.v * (255 - ((hsv.s * remainder) >> 8))) >> 8;
    t = (hsv.v * (255 - ((hsv.s * (255 - remainder)) >> 8))) >> 8;

    switch (region)
    {
    case 0:
        rgb.r = hsv.v;
        rgb.g = t;
        rgb.b = p;
        break;
    case 1:
        rgb.r = q;
        rgb.g = hsv.v;
        rgb.b = p;
        break;
    case 2:
        rgb.r = p;
        rgb.g = hsv.v;
        rgb.b = t;
        break;
    case 3:
        rgb.r = p;
        rgb.g = q;
        rgb.b = hsv.v;
        break;
    case 4:
        rgb.r = t;
        rgb.g = p;
        rgb.b = hsv.v;
        break;
    default:
        rgb.r = hsv.v;
        rgb.g = p;
        rgb.b = q;
        break;
    }

    return rgb;
}

uint16_t led_get_num_leds_a(void) { return strip_a.num_leds; }
uint16_t led_get_num_leds_b(void) { return strip_b.num_leds; }

rgb_t led_get_pixel_a(uint16_t index)
{
    rgb_t color = {0};
    if (!strip_a.initialized || index >= strip_a.num_leds)
        return color;
    xSemaphoreTake(led_mutex, portMAX_DELAY);
    color = strip_a.buffer[index];
    xSemaphoreGive(led_mutex);
    return color;
}

rgb_t led_get_pixel_b(uint16_t index)
{
    rgb_t color = {0};
    if (!strip_b.initialized || index >= strip_b.num_leds)
        return color;
    xSemaphoreTake(led_mutex, portMAX_DELAY);
    color = strip_b.buffer[index];
    xSemaphoreGive(led_mutex);
    return color;
}

void led_add_pixel_a(uint16_t index, rgb_t color)
{
    if (!strip_a.initialized || index >= strip_a.num_leds)
        return;
    xSemaphoreTake(led_mutex, portMAX_DELAY);
    strip_a.buffer[index].r = (strip_a.buffer[index].r + color.r > 255) ? 255 : strip_a.buffer[index].r + color.r;
    strip_a.buffer[index].g = (strip_a.buffer[index].g + color.g > 255) ? 255 : strip_a.buffer[index].g + color.g;
    strip_a.buffer[index].b = (strip_a.buffer[index].b + color.b > 255) ? 255 : strip_a.buffer[index].b + color.b;
    xSemaphoreGive(led_mutex);
}

void led_add_pixel_b(uint16_t index, rgb_t color)
{
    if (!strip_b.initialized || index >= strip_b.num_leds)
        return;
    xSemaphoreTake(led_mutex, portMAX_DELAY);
    strip_b.buffer[index].r = (strip_b.buffer[index].r + color.r > 255) ? 255 : strip_b.buffer[index].r + color.r;
    strip_b.buffer[index].g = (strip_b.buffer[index].g + color.g > 255) ? 255 : strip_b.buffer[index].g + color.g;
    strip_b.buffer[index].b = (strip_b.buffer[index].b + color.b > 255) ? 255 : strip_b.buffer[index].b + color.b;
    xSemaphoreGive(led_mutex);
}