/* WS2812 driver — from https://github.com/ForsakenNGS/Pico_WS2812 (BSD-style) */

#include "WS2812.hpp"

#include "WS2812.pio.h"

#include "hardware/clocks.h"
#include "hardware/gpio.h"

#include <cstdlib>

WS2812::WS2812(uint pin, uint length, PIO pio, uint sm) {
    initialize(pin, length, pio, sm, NONE, GREEN, RED, BLUE);
}

WS2812::WS2812(uint pin, uint length, PIO pio, uint sm, DataFormat format) {
    switch (format) {
    case FORMAT_RGB:
        initialize(pin, length, pio, sm, NONE, RED, GREEN, BLUE);
        break;
    case FORMAT_GRB:
        initialize(pin, length, pio, sm, NONE, GREEN, RED, BLUE);
        break;
    case FORMAT_WRGB:
        initialize(pin, length, pio, sm, WHITE, RED, GREEN, BLUE);
        break;
    }
}

WS2812::~WS2812() {
    delete[] data;
}

void WS2812::initialize(uint pin, uint length, PIO pio, uint sm, DataByte b1, DataByte b2,
                        DataByte b3, DataByte b4) {
    this->pin = pin;
    this->length = length;
    this->pio = pio;
    this->sm = sm;
    this->data = new uint32_t[length];
    this->bytes[0] = b1;
    this->bytes[1] = b2;
    this->bytes[2] = b3;
    this->bytes[3] = b4;
    this->bits = (b1 == NONE ? 24 : 32);
    /* One copy of the program per PIO block — shared by all strips on that PIO. */
    static int pio0_offset = -1;
    static int pio1_offset = -1;
    int *cached = (pio == pio0) ? &pio0_offset : &pio1_offset;
    if (*cached < 0) {
        *cached = (int)pio_add_program(pio, &ws2812_program);
    }
    this->offset = (uint)*cached;
    applyPinConfig();
}

void WS2812::applyPinConfig() {
    pio_sm_set_enabled(pio, sm, false);
    pio_gpio_init(pio, pin);
    pio_sm_set_consecutive_pindirs(pio, sm, pin, 1, true);

    pio_sm_config c = ws2812_program_get_default_config(offset);
    sm_config_set_sideset_pins(&c, pin);
    sm_config_set_out_shift(&c, false, true, bits);
    sm_config_set_fifo_join(&c, PIO_FIFO_JOIN_TX);

    int cycles_per_bit = ws2812_T1 + ws2812_T2 + ws2812_T3;
    float div = (float)clock_get_hz(clk_sys) / (800000.f * (float)cycles_per_bit);
    sm_config_set_clkdiv(&c, div);

    pio_sm_init(pio, sm, offset, &c);
    pio_sm_set_enabled(pio, sm, true);
}

void WS2812::setPin(uint new_pin) {
    if (new_pin == pin) {
        return;
    }
    pio_sm_set_enabled(pio, sm, false);
    gpio_set_function(pin, GPIO_FUNC_SIO);
    gpio_set_dir(pin, GPIO_IN);
    pin = new_pin;
    applyPinConfig();
}

uint32_t WS2812::convertData(uint32_t rgbw) {
    uint32_t result = 0;
    for (uint b = 0; b < 4; b++) {
        switch (bytes[b]) {
        case RED:
            result |= (rgbw & 0xFF);
            break;
        case GREEN:
            result |= (rgbw & 0xFF00) >> 8;
            break;
        case BLUE:
            result |= (rgbw & 0xFF0000) >> 16;
            break;
        case WHITE:
            result |= (rgbw & 0xFF000000) >> 24;
            break;
        default:
            break;
        }
        result <<= 8;
    }
    return result;
}

void WS2812::setPixelColor(uint index, uint32_t color) {
    if (index < length) {
        data[index] = convertData(color);
    }
}

void WS2812::setPixelColor(uint index, uint8_t red, uint8_t green, uint8_t blue) {
    setPixelColor(index, RGB(red, green, blue));
}

void WS2812::fill(uint32_t color) {
    color = convertData(color);
    for (uint i = 0; i < length; i++) {
        data[i] = color;
    }
}

void WS2812::show() {
    show(length);
}

void WS2812::show(uint count) {
    if (count > length) {
        count = length;
    }
    for (uint i = 0; i < count; i++) {
        pio_sm_put_blocking(pio, sm, data[i]);
    }
}