#include <SystemK.h>
#include <freertos/FreeRTOS.h>
#include <freertos/task.h>
#include <string.h>
#include <muxed_led_strip.h>
#include <muxed_led_strip_interface.h>
#include <driver/gpio.h>
#include <driver/rmt_tx.h>
#include "HW_NeoPixels.h"

#define NEOPIXELS_STACK_SIZE (100 * 1024)
#define NEOPIXELS_TASK_PRIORITY (tskIDLE_PRIORITY + 1)
static StaticTask_t xTaskBuffer;
StackType_t *xStack;
static TaskHandle_t xTaskHandle;

static muxed_led_strip_handle_t NeoPixel_Out;

// GPIO assignments
#define NEOPIXEL_OUT_GPIO      GPIO_NUM_48
#define BARREL_ENABLE_GPIO     GPIO_NUM_18
#define RECEIVER_ENABLE_GPIO   GPIO_NUM_45
#define DISPLAY_ENABLE_GPIO    GPIO_NUM_10
#define EFFECTS_ENABLE_GPIO    GPIO_NUM_17

// 10MHz resolution, 1 tick = 0.1us (led strip needs a high resolution)
#define LED_STRIP_RMT_RES_HZ (10 * 1000 * 1000)

static const char *TAG = "NeoPixels";

static inline void Enable_Channel(NeoPixelsChannel_T channel)
{
    // The Enable lines are active low.
    if (channel == NEOPIXEL_CHANNEL_ALL)
    {
        gpio_set_level(BARREL_ENABLE_GPIO, 0);
        gpio_set_level(RECEIVER_ENABLE_GPIO, 0);
        gpio_set_level(DISPLAY_ENABLE_GPIO, 0);
        gpio_set_level(EFFECTS_ENABLE_GPIO, 0);
    }
    else if (channel < CONFIG_KTAG_N_NEOPIXEL_CHANNELS)
    {
        gpio_set_level(BARREL_ENABLE_GPIO, channel == NEOPIXEL_CHANNEL_BARREL ? 0 : 1);
        gpio_set_level(RECEIVER_ENABLE_GPIO, channel == NEOPIXEL_CHANNEL_RECEIVER ? 0 : 1);
        gpio_set_level(DISPLAY_ENABLE_GPIO, channel == NEOPIXEL_CHANNEL_DISPLAY ? 0 : 1);
        gpio_set_level(EFFECTS_ENABLE_GPIO, channel == NEOPIXEL_CHANNEL_EFFECTS ? 0 : 1);
    }
    else
    {
        // Select none.
        gpio_set_level(BARREL_ENABLE_GPIO, 1);
        gpio_set_level(RECEIVER_ENABLE_GPIO, 1);
        gpio_set_level(DISPLAY_ENABLE_GPIO, 1);
        gpio_set_level(EFFECTS_ENABLE_GPIO, 1);
    }
}

static bool IRAM_ATTR rmt_tx_done_callback(rmt_channel_handle_t channel, const rmt_tx_done_event_data_t *edata, void *user_data)
{
    BaseType_t xHigherPriorityTaskWoken = pdFALSE;

    // Send a notification to the task that the current channel is complete.
    xTaskNotifyIndexedFromISR(xTaskHandle,
                              0,
                              0,
                              eSetBits,
                              &xHigherPriorityTaskWoken);

    return xHigherPriorityTaskWoken;
}

static rmt_tx_event_callbacks_t cbs = {
    .on_trans_done = rmt_tx_done_callback,
};

static gpio_config_t barrel_enable_gpio_config = {
    .pin_bit_mask = (1ULL << BARREL_ENABLE_GPIO),
    .mode = GPIO_MODE_OUTPUT, // Set mode to output
    .pull_up_en = GPIO_PULLUP_ENABLE,
    .pull_down_en = GPIO_PULLDOWN_DISABLE,
    .intr_type = GPIO_INTR_DISABLE};

static gpio_config_t receiver_enable_gpio_config = {
    .pin_bit_mask = (1ULL << RECEIVER_ENABLE_GPIO),
    .mode = GPIO_MODE_OUTPUT, // Set mode to output
    .pull_up_en = GPIO_PULLUP_ENABLE,
    .pull_down_en = GPIO_PULLDOWN_DISABLE,
    .intr_type = GPIO_INTR_DISABLE};

static gpio_config_t display_enable_gpio_config = {
    .pin_bit_mask = (1ULL << DISPLAY_ENABLE_GPIO),
    .mode = GPIO_MODE_OUTPUT, // Set mode to output
    .pull_up_en = GPIO_PULLUP_ENABLE,
    .pull_down_en = GPIO_PULLDOWN_DISABLE,
    .intr_type = GPIO_INTR_DISABLE};

static gpio_config_t effects_enable_gpio_config = {
    .pin_bit_mask = (1ULL << EFFECTS_ENABLE_GPIO),
    .mode = GPIO_MODE_OUTPUT, // Set mode to output
    .pull_up_en = GPIO_PULLUP_ENABLE,
    .pull_down_en = GPIO_PULLDOWN_DISABLE,
    .intr_type = GPIO_INTR_DISABLE};

static muxed_led_strip_config_t neopixel_out_config = {
    .strip_gpio_num = NEOPIXEL_OUT_GPIO,               // The GPIO that connected to the LED strip's data line
    .channels = CONFIG_KTAG_N_NEOPIXEL_CHANNELS,       // The number of multiplexed channels
    .max_leds = CONFIG_KTAG_MAX_NEOPIXELS_PER_CHANNEL, // The number of LEDs in the strip
    .led_pixel_format = LED_PIXEL_FORMAT_GRB,          // Pixel format of your LED strip
    .led_model = LED_MODEL_WS2812,                     // LED strip model
    .flags.invert_out = false,                         // whether to invert the output signal
};

// From <https://www.espressif.com/sites/default/files/documentation/esp32-s3_technical_reference_manual_en.pdf#rmt>:
//
// The RMT module has eight channels, numbered from zero to seven. Each channel is able to independently
// transmit or receive signals.
//    • Channel 0 ~ 3 (TX channel) are dedicated to sending signals.
//    • Channel 4 ~ 7 (RX channel) are dedicated to receiving signals.
// Each TX/RX channel is controlled by a dedicated set of registers with the same functionality. Channel 3 and
// channel 7 support DMA access, so the two channels also have a set of DMA-related control and status registers

static led_strip_rmt_config_t rmt_config_noDMA = {
    .clk_src = RMT_CLK_SRC_XTAL,           // different clock source can lead to different power consumption
    .resolution_hz = LED_STRIP_RMT_RES_HZ, // RMT counter clock frequency
    .flags.with_dma = false,               // Only the last channel has the DMA capability
};

static led_strip_spi_config_t spi_config_withDMA __attribute__((unused)) = {
    .clk_src = RMT_CLK_SRC_XTAL,    // different clock source can lead to different power consumption
    .flags.with_dma = true,         // Using DMA can improve performance and help drive more LEDs
    .spi_bus = SPI2_HOST,           // SPI bus ID
};

void Initialize_SystemK_NeoPixels(SemaphoreHandle_t init_complete)
{
    xStack = (uint8_t *)heap_caps_calloc(1, NEOPIXELS_STACK_SIZE, MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT | MALLOC_CAP_32BIT);

    // Create the task without using any dynamic memory allocation.
    xTaskHandle = xTaskCreateStaticPinnedToCore(
        NeoPixels_Task,          // Function that implements the task (this is part of SystemK).
        "NeoPixels",             // Text name for the task.
        NEOPIXELS_STACK_SIZE,    // Number of indexes in the xStack array.
        NULL,                    // Parameter passed into the task.
        NEOPIXELS_TASK_PRIORITY, // Priority at which the task is created.
        xStack,                  // Array to use as the task's stack.
        &xTaskBuffer,            // Variable to hold the task's data structure.
        APP_CPU_NUM);            // Specify the task's core affinity.

    KLOG_INFO(TAG, "Initialization complete.");

    xSemaphoreGive(init_complete);
}

SystemKResult_T HW_NeoPixels_Init(void)
{
    if (CONFIG_KTAG_N_NEOPIXEL_CHANNELS > 0)
    {
        // Initialize the NeoPixel Out and the Barrel Enable.
        ESP_ERROR_CHECK(muxed_led_strip_new_rmt_device(&neopixel_out_config, &rmt_config_noDMA, &cbs, &NeoPixel_Out));
        KLOG_INFO(TAG, "Initialized NeoPixel Out as GPIO[%d].", NEOPIXEL_OUT_GPIO);

        ESP_ERROR_CHECK(gpio_config(&barrel_enable_gpio_config));
        KLOG_INFO(TAG, "Initialized barrel NeoPixel enable as GPIO[%d].", BARREL_ENABLE_GPIO);
    }
    if (CONFIG_KTAG_N_NEOPIXEL_CHANNELS > 1)
    {
        ESP_ERROR_CHECK(gpio_config(&receiver_enable_gpio_config));
        KLOG_INFO(TAG, "Initialized receiver NeoPixel enable as GPIO[%d].", RECEIVER_ENABLE_GPIO);
    }
    if (CONFIG_KTAG_N_NEOPIXEL_CHANNELS > 2)
    {
        ESP_ERROR_CHECK(gpio_config(&display_enable_gpio_config));
        KLOG_INFO(TAG, "Initialized display NeoPixel enable as GPIO[%d].", DISPLAY_ENABLE_GPIO);
    }
    if (CONFIG_KTAG_N_NEOPIXEL_CHANNELS > 3)
    {
        ESP_ERROR_CHECK(gpio_config(&effects_enable_gpio_config));
        KLOG_INFO(TAG, "Initialized effects NeoPixel enable as GPIO[%d].", EFFECTS_ENABLE_GPIO);
    }

    return SYSTEMK_RESULT_SUCCESS;
}

SystemKResult_T HW_NeoPixels_Set_RGB(NeoPixelsChannel_T channel, uint8_t position, uint8_t red, uint8_t green, uint8_t blue)
{
    if (channel < CONFIG_KTAG_N_NEOPIXEL_CHANNELS)
    {
        // TODO: Add code to account for RGB order in each channel.
        ESP_ERROR_CHECK(muxed_led_strip_set_pixel(NeoPixel_Out, channel, position, red, green, blue));
    }

    return SYSTEMK_RESULT_SUCCESS;
}

SystemKResult_T HW_NeoPixels_Publish(void)
{
    for (uint_fast8_t channel = 0; channel < CONFIG_KTAG_N_NEOPIXEL_CHANNELS; channel++)
    {
        uint32_t channel_complete;
        Enable_Channel(channel);
        ESP_ERROR_CHECK_WITHOUT_ABORT(muxed_led_strip_refresh(NeoPixel_Out, channel));
        xTaskNotifyWaitIndexed(0,                 /* Wait for 0th Notification */
                               0x00,              /* Don't clear any bits on entry. */
                               UINT32_MAX,        /* Clear all bits on exit. */
                               &channel_complete, /* Receives the notification value. */
                               portMAX_DELAY);    /* Block indefinitely. */
        // KLOG_INFO(TAG, "Published %lu.", channel_complete);
    }
    Enable_Channel(NEOPIXEL_CHANNEL_NONE);

    return SYSTEMK_RESULT_SUCCESS;
}

color_t HW_NeoPixels_Get_My_Color(void)
{
    color_t result = COLOR_ORANGE;

    uint8_t Team_ID;
    uint8_t Player_ID;
    uint8_t Weapon_ID;

    (void) SETTINGS_get_uint8_t(SYSTEMK_SETTING_TEAMID, &Team_ID);
    (void) SETTINGS_get_uint8_t(SYSTEMK_SETTING_PLAYERID, &Player_ID);
    (void) SETTINGS_get_uint8_t(SYSTEMK_SETTING_WEAPONID, &Weapon_ID);

    result = PROTOCOLS_GetColor(GetWeaponFromID(Weapon_ID).Protocol, Team_ID, Player_ID); 

    return result;
}