#include <zephyr/kernel.h>
#include <nrfx_pwm.h>
#include <hal/nrf_gpio.h>

#define PWM_INST_IDX 21
#define MAX_PWM_INTERFACES  1

//PWM PIN Define
#define PWM_OUTPUT_PIN1 NRF_GPIO_PIN_MAP(1, 10) /* Replace with your desired GPIO pin */
#define PWM_OUTPUT_PIN2 NRF_GPIO_PIN_MAP(1, 11) /* Replace with your desired GPIO pin */
#define PWM_OUTPUT_PIN3 NRF_GPIO_PIN_MAP(1, 12) /* Replace with your desired GPIO pin */
#define PWM_OUTPUT_PIN4 NRF_GPIO_PIN_MAP(1, 13) /* Replace with your desired GPIO pin */

// Define PWM instance
static nrfx_pwm_t pwm_instance = NRFX_PWM_INSTANCE(PWM_INST_IDX);

#define DIGITAL_BYTES_PER_WRITE     24
#define MULTIPLY_FACTOR             8

// below will be going into neostrip.c
#define NO_OF_LEDS_PER_CHANNEL       24
#define NO_OF_CHANNELS               4
#define BYTES_PER_LED                5
#define TOTAL_LEDS                   (NO_OF_LEDS_PER_CHANNEL * NO_OF_CHANNELS)
static uint8_t byte_buf[TOTAL_LEDS * BYTES_PER_LED];

static uint8_t multiply_factor = 0;

typedef enum  __PACKED_ENUM_TYPE__ {
    LED_STRIP_IC_BULLET_PIXEL = 0x4C,        // STRIP_TYPE_C9, STRIP_TYPE_BULLET, STRIP_TYPE_NUGGET
} strip_ic_type;

static uint8_t color[][3] = {   {0x00, 0x00, 0xFF},
                                {0xFF, 0x00, 0x00},
                                {0x00, 0xFF, 0x00} };

uint8_t ic_type = LED_STRIP_IC_BULLET_PIXEL;

// Sequence buffers (must be in RAM)
static nrf_pwm_values_individual_t seq0_values[DIGITAL_BYTES_PER_WRITE * MULTIPLY_FACTOR] = {{10 | 0x8000}, {10 | 0x8000}, {10 | 0x8000}, {10 | 0x8000}, {10 | 0x8000}, {10 | 0x8000}, {10 | 0x8000}, {10 | 0x8000}, {10 | 0x8000}, {10 | 0x8000}, {10 | 0x8000}, {10 | 0x8000}, {10 | 0x8000}, {10 | 0x8000}, {10 | 0x8000}, {10 | 0x8000}, {10 | 0x8000}, {10 | 0x8000}, {10 | 0x8000}, {10 | 0x8000}, {10 | 0x8000}, {10 | 0x8000}, {10 | 0x8000}, {10 | 0x8000}}; // Initial duty cycles for SEQ0
static nrf_pwm_values_individual_t seq1_values[DIGITAL_BYTES_PER_WRITE * MULTIPLY_FACTOR] = {{5 | 0x8000}, {5 | 0x8000}, {5 | 0x8000}, {5 | 0x8000}, {5 | 0x8000}, {5 | 0x8000}, {5 | 0x8000}, {5 | 0x8000}, {5 | 0x8000}, {5 | 0x8000}, {5 | 0x8000}, {5 | 0x8000}, {5 | 0x8000}, {5 | 0x8000}, {5 | 0x8000}, {5 | 0x8000}, {5 | 0x8000}, {5 | 0x8000}, {5 | 0x8000}, {5 | 0x8000}, {5 | 0x8000}, {5 | 0x8000}, {5 | 0x8000}, {5 | 0x8000}}; // Initial duty cycles for SEQ1

typedef struct {
    uint16_t pwm_cycle_count; //total pwm cycle count
    uint16_t bit1_count; //pwm cycle for bit 1
    uint16_t bit0_count; //pwm cycle for bit 0
    nrf_pwm_clk_t pwm_clk_freq; //pwm clock frequency
} LED_Driver_property_t;

typedef struct {
    uint8_t *buf0;
    uint16_t pending_bytes0;
    uint8_t *buf1;
    uint16_t pending_bytes1;
    uint8_t *buf2;
    uint16_t pending_bytes2;
    uint8_t *buf3;
    uint16_t pending_bytes3;
    LED_Driver_property_t * property;
    uint32_t bit_seq[4];
    nrf_pwm_values_individual_t *seq_ind_buf[2];            // size of [2][PWM_SEQ_AMOUNT] ping pong, PWM_SEQ_AMOUNT must power of 2
    nrf_pwm_sequence_t seq1;
    nrf_pwm_sequence_t seq2;
    strip_ic_type ic_type;
    uint8_t is_pending;
} PWM_TX_Status_t;

PWM_TX_Status_t pwm_handle[MAX_PWM_INTERFACES];

const LED_Driver_property_t ws2812b = {
        .pwm_cycle_count = 20,
        .bit1_count = 15,
        .bit0_count = 5,
        .pwm_clk_freq = NRF_PWM_CLK_16MHz,
};

/* PWM Functions Implementation starts */

static void update_multi_strip_dma_pwm(uint8_t instance, uint8_t *buf0, uint16_t count0,
                                                        uint8_t *buf1, uint16_t count1,
                                                        uint8_t *buf2, uint16_t count2,
                                                        uint8_t *buf3, uint16_t count3,
                                                        uint32_t *seq_ptr)
{
    uint8_t b;
    nrf_pwm_values_individual_t *seq_ind_buf = (nrf_pwm_values_individual_t *)seq_ptr;

    for (uint8_t k = 0; k < DIGITAL_BYTES_PER_WRITE; k++) 
    {   //don't loop from 0 to 3 to increase spped
        // byte value as aa-bb-cc-dd
        // map dd to [3], map cc to [2], map bb to [1], map aa to [0]

        if (k < count0)
            b = buf0[k];
        else
            b = 0; 
        seq_ind_buf[7].channel_0 = (pwm_handle[instance].bit_seq[b & 0x3] >> 16); 
        seq_ind_buf[6].channel_0 = (pwm_handle[instance].bit_seq[b & 0x3] & 0xFFFF);
        b >>= 2;
        seq_ind_buf[5].channel_0 = (pwm_handle[instance].bit_seq[b & 0x3] >> 16); 
        seq_ind_buf[4].channel_0 = (pwm_handle[instance].bit_seq[b & 0x3] & 0xFFFF);
        b >>= 2;
        seq_ind_buf[3].channel_0 = (pwm_handle[instance].bit_seq[b & 0x3] >> 16); 
        seq_ind_buf[2].channel_0 = (pwm_handle[instance].bit_seq[b & 0x3] & 0xFFFF);
        b >>= 2;
        seq_ind_buf[1].channel_0 = (pwm_handle[instance].bit_seq[b & 0x3] >> 16); 
        seq_ind_buf[0].channel_0 = (pwm_handle[instance].bit_seq[b & 0x3] & 0xFFFF);

        if (k < count1)
            b = buf1[k];
        else
            b = 0;
        seq_ind_buf[7].channel_1 = (pwm_handle[instance].bit_seq[b & 0x3] >> 16);
        seq_ind_buf[6].channel_1 = (pwm_handle[instance].bit_seq[b & 0x3] & 0xFFFF);
        b >>= 2;
        seq_ind_buf[5].channel_1 = (pwm_handle[instance].bit_seq[b & 0x3] >> 16);
        seq_ind_buf[4].channel_1 = (pwm_handle[instance].bit_seq[b & 0x3] & 0xFFFF);
        b >>= 2;
        seq_ind_buf[3].channel_1 = (pwm_handle[instance].bit_seq[b & 0x3] >> 16);
        seq_ind_buf[2].channel_1 = (pwm_handle[instance].bit_seq[b & 0x3] & 0xFFFF);
        b >>= 2;
        seq_ind_buf[1].channel_1 = (pwm_handle[instance].bit_seq[b & 0x3] >> 16);
        seq_ind_buf[0].channel_1 = (pwm_handle[instance].bit_seq[b & 0x3] & 0xFFFF);

        if (k < count2)
            b = buf2[k];
        else
            b = 0;
        seq_ind_buf[7].channel_2 = (pwm_handle[instance].bit_seq[b & 0x3] >> 16);
        seq_ind_buf[6].channel_2 = (pwm_handle[instance].bit_seq[b & 0x3] & 0xFFFF);
        b >>= 2;
        seq_ind_buf[5].channel_2 = (pwm_handle[instance].bit_seq[b & 0x3] >> 16);
        seq_ind_buf[4].channel_2 = (pwm_handle[instance].bit_seq[b & 0x3] & 0xFFFF);
        b >>= 2;
        seq_ind_buf[3].channel_2 = (pwm_handle[instance].bit_seq[b & 0x3] >> 16);
        seq_ind_buf[2].channel_2 = (pwm_handle[instance].bit_seq[b & 0x3] & 0xFFFF);
        b >>= 2;
        seq_ind_buf[1].channel_2 = (pwm_handle[instance].bit_seq[b & 0x3] >> 16);
        seq_ind_buf[0].channel_2 = (pwm_handle[instance].bit_seq[b & 0x3] & 0xFFFF);

        if (k < count3)
            b = buf3[k];
        else
            b = 0;
        seq_ind_buf[7].channel_3 = (pwm_handle[instance].bit_seq[b & 0x3] >> 16);
        seq_ind_buf[6].channel_3 = (pwm_handle[instance].bit_seq[b & 0x3] & 0xFFFF);
        b >>= 2;
        seq_ind_buf[5].channel_3 = (pwm_handle[instance].bit_seq[b & 0x3] >> 16);
        seq_ind_buf[4].channel_3 = (pwm_handle[instance].bit_seq[b & 0x3] & 0xFFFF);
        b >>= 2;
        seq_ind_buf[3].channel_3 = (pwm_handle[instance].bit_seq[b & 0x3] >> 16);
        seq_ind_buf[2].channel_3 = (pwm_handle[instance].bit_seq[b & 0x3] & 0xFFFF);
        b >>= 2;
        seq_ind_buf[1].channel_3 = (pwm_handle[instance].bit_seq[b & 0x3] >> 16);
        seq_ind_buf[0].channel_3 = (pwm_handle[instance].bit_seq[b & 0x3] & 0xFFFF);
        seq_ind_buf += 8;
    }
}


void update_sequence(uint8_t instance, uint8_t seq_id)
{
    uint16_t minimum_bytes = DIGITAL_BYTES_PER_WRITE;

    uint16_t c0 = MIN(pwm_handle[instance].pending_bytes0, minimum_bytes);
    uint16_t c1 = MIN(pwm_handle[instance].pending_bytes1, minimum_bytes);
    uint16_t c2 = MIN(pwm_handle[instance].pending_bytes2, minimum_bytes);
    uint16_t c3 = MIN(pwm_handle[instance].pending_bytes3, minimum_bytes);
    uint32_t *p = (uint32_t *) pwm_handle[instance].seq_ind_buf[seq_id];

    // printf("updating seq id: %d bytes: %d\n", seq_id, c0);
    update_multi_strip_dma_pwm(instance, pwm_handle[instance].buf0, c0,
                            pwm_handle[instance].buf1, c1,
                            pwm_handle[instance].buf2, c2,
                            pwm_handle[instance].buf3, c3, p);

    pwm_handle[instance].buf0 += c0;
    pwm_handle[instance].pending_bytes0 -= c0;
    pwm_handle[instance].buf1 += c1;
    pwm_handle[instance].pending_bytes1 -= c1;
    pwm_handle[instance].buf2 += c2;
    pwm_handle[instance].pending_bytes2 -= c2;
    pwm_handle[instance].buf3 += c3;
    pwm_handle[instance].pending_bytes3 -= c3;
}

// Event handler
static void pwm_handler(nrfx_pwm_evt_type_t event_type, void *p_context)
{
    switch (event_type) {
        case NRFX_PWM_EVT_END_SEQ0:
            // SEQ0 finished, safe to update SEQ0 values
            // printf("pwm contex %d\n", ((nrfx_pwm_t *)p_context)->instance_id);
            update_sequence(0, 0);
            nrfx_pwm_sequence_update(&pwm_instance, 0,  &(pwm_handle[0].seq1));
            // printk("SEQ0 updated\n");
            break;
        case NRFX_PWM_EVT_END_SEQ1:
            // SEQ1 finished, safe to update SEQ1 if needed
            update_sequence(0, 1);
            nrfx_pwm_sequence_update(&pwm_instance, 1,  &(pwm_handle[0].seq2));
            // printk("SEQ1 updated\n");
            break;
        case NRFX_PWM_EVT_FINISHED:
            // printk("Playback finished\n");
            break;
        case NRFX_PWM_EVT_STOPPED:
            // printk("PWM stopped\n");
            break;
        default:
            break;
    }
}

uint8_t multi_strip_digital_init(uint8_t instance, uint8_t type)
{
    uint16_t bit0val, bit1val;
    nrfx_err_t err;
    printf("Init IC: %d\r\n", type);
    if (type == LED_STRIP_IC_BULLET_PIXEL)
    {
        pwm_handle[instance].property = (LED_Driver_property_t*) &ws2812b;

        bit0val = 0x8000 | pwm_handle[instance].property->bit0_count;
        bit1val = 0x8000 | pwm_handle[instance].property->bit1_count;
    }

    pwm_handle[instance].bit_seq[0] = (bit0val <<16) | bit0val; // "0" at lower 16bit since it's sent out first
    pwm_handle[instance].bit_seq[1] = (bit1val <<16) | bit0val;
    pwm_handle[instance].bit_seq[2] = (bit0val <<16) | bit1val;
    pwm_handle[instance].bit_seq[3] = (bit1val <<16) | bit1val;

    pwm_handle[instance].is_pending = 0;
    pwm_handle[instance].seq_ind_buf[0] = seq0_values;
    pwm_handle[instance].seq_ind_buf[1] = seq1_values;

    memset(pwm_handle[instance].seq_ind_buf[0], 0, sizeof(seq0_values));
    memset(pwm_handle[instance].seq_ind_buf[1], 0, sizeof(seq1_values));

    pwm_handle[instance].seq1.values.p_individual = pwm_handle[instance].seq_ind_buf[0];
    pwm_handle[instance].seq1.length =  NRF_PWM_VALUES_LENGTH(seq0_values);
    pwm_handle[instance].seq1.repeats = 0;

    pwm_handle[instance].seq2.values.p_individual = pwm_handle[instance].seq_ind_buf[1];
    pwm_handle[instance].seq2.length =  NRF_PWM_VALUES_LENGTH(seq1_values);
    pwm_handle[instance].seq2.repeats = 0;

    if(instance == 0)
    {
        nrfx_pwm_config_t pwm_config = {
            .output_pins = {
                PWM_OUTPUT_PIN1,  // channel 1
                PWM_OUTPUT_PIN2,  // channel 2
                PWM_OUTPUT_PIN3,   // channel 3
                PWM_OUTPUT_PIN4,  // channel 0
            },
            .base_clock = pwm_handle[instance].property->pwm_clk_freq,
            .count_mode = NRF_PWM_MODE_UP,
            .top_value  = pwm_handle[instance].property->pwm_cycle_count,
            .load_mode  = NRF_PWM_LOAD_INDIVIDUAL,
            .step_mode  = NRF_PWM_STEP_AUTO,
            .irq_priority = NRFX_PWM_DEFAULT_CONFIG_IRQ_PRIORITY,
        };

        err = nrfx_pwm_init(&pwm_instance, &pwm_config, pwm_handler, &pwm_instance);
        if (err != NRFX_SUCCESS) {
            printf("Failed to initialize PWM: %d\n", err);
            return 0;
        }

        printf("PWM initialized successfully.\n");
    }
    return 0;
}

uint8_t led_strip_pwm_interface_init(uint8_t ic_type)
{
        uint8_t instance = 0;
        pwm_handle[instance].ic_type = ic_type;
        multi_strip_digital_init(instance, ic_type);
        return 0;
}

void led_multi_strip_pwm_interface_transfer(uint8_t ic_type)
{
    uint16_t byte_cnt0, byte_cnt1, byte_cnt2, byte_cnt3;
    uint8_t instance = 0;
    uint16_t maxbyteCnt = 0;
    uint32_t flags = 0x00, ret = 0;

    int loop_times = 1;
    switch(ic_type){
        case LED_STRIP_IC_BULLET_PIXEL:
            //TODO: This will be updated as per no of leds configured
            byte_cnt0 = byte_cnt1 = byte_cnt2 = byte_cnt3 = (NO_OF_LEDS_PER_CHANNEL * multiply_factor);
            maxbyteCnt = byte_cnt0;

            pwm_handle[instance].is_pending = 1;
            pwm_handle[instance].buf0 = &byte_buf[(NO_OF_LEDS_PER_CHANNEL * multiply_factor) * 0];
            pwm_handle[instance].pending_bytes0 = byte_cnt0;
            pwm_handle[instance].buf1 = &byte_buf[(NO_OF_LEDS_PER_CHANNEL * multiply_factor) * 1];
            pwm_handle[instance].pending_bytes1 = byte_cnt1;
            pwm_handle[instance].buf2 = &byte_buf[(NO_OF_LEDS_PER_CHANNEL * multiply_factor) * 2];
            pwm_handle[instance].pending_bytes2 = byte_cnt2;
            pwm_handle[instance].buf3 = &byte_buf[(NO_OF_LEDS_PER_CHANNEL * multiply_factor) * 3];
            pwm_handle[instance].pending_bytes3 = byte_cnt3;

            printf("bytebuf: 0x%X size: %d\r\n", byte_buf, sizeof(byte_buf));
            printf("buf0: 0x%X buf1: 0x%X buf3: 0x%X buf4: 0x%X\r\n", 
                                                pwm_handle[instance].buf0,
                                                pwm_handle[instance].buf1,
                                                pwm_handle[instance].buf2,
                                                pwm_handle[instance].buf3);
            printf("adding seq 0\r\n");
            update_sequence(instance, 0);

            if ((pwm_handle[instance].pending_bytes0 > 0)
                || (pwm_handle[instance].pending_bytes1 > 0)
                || (pwm_handle[instance].pending_bytes2 > 0)
                || (pwm_handle[instance].pending_bytes3 > 0)) {
                printf("adding seq 1\r\n");
                update_sequence(instance, 1);
            }

            loop_times = maxbyteCnt / (2 * DIGITAL_BYTES_PER_WRITE);
            if (maxbyteCnt % (2 * DIGITAL_BYTES_PER_WRITE))
                loop_times += 1;

            printf("loop times: %d\n", loop_times);
            flags = NRFX_PWM_FLAG_STOP | NRFX_PWM_FLAG_SIGNAL_END_SEQ0 | NRFX_PWM_FLAG_SIGNAL_END_SEQ1;
            ret = nrfx_pwm_complex_playback(&pwm_instance, &pwm_handle[instance].seq1, \
                &pwm_handle[instance].seq2, loop_times, flags);
            printf("return value: %d\n", ret);
            printk("PWM playback started\n");
            break;
    }
}

void copy_led_to_buf(uint8_t ic_type)
{
    switch(ic_type)
    {
        case LED_STRIP_IC_BULLET_PIXEL:
        {
            int i = 0, j = 0;
            multiply_factor = 3;
            for(int i = 0; i < TOTAL_LEDS; i++)
            {
                byte_buf[j++] = color[i % 3][0];
                byte_buf[j++] = color[i % 3][1];
                byte_buf[j++] = color[i % 3][2];
                printf("led: %d color: %d %d %d\r\n",i , byte_buf[j - 3], byte_buf[j - 2], byte_buf[j - 1]);
            }
        }
        break;
    }
}

void main(void)
{
    uint32_t flags = 0x00;
    int err;

#if defined(__ZEPHYR__)
        IRQ_CONNECT(NRFX_IRQ_NUMBER_GET(NRF_PWM_INST_GET(PWM_INST_IDX)), IRQ_PRIO_LOWEST,
                NRFX_PWM_INST_HANDLER_GET(PWM_INST_IDX), 0, 0);
#endif
    led_strip_pwm_interface_init(ic_type);

    copy_led_to_buf(ic_type);
    led_multi_strip_pwm_interface_transfer(ic_type);
    // Start complex playback (play SEQ0 then SEQ1, repeat 2 times, then stop)

    // Keep the system running
    while (1) {
        k_sleep(K_MSEC(1000));
    }
}
