/* Multi-strip WS2812 wrapper around ForsakenNGS/Pico_WS2812 */

#include "ws2812_led.h"

#include "board_config.h"
#include "WS2812.hpp"

#include "hardware/pio.h"
#include "pico/stdlib.h"

typedef struct {
    WS2812 *strip;
    PIO pio;
    uint sm;
    unsigned int pin;
    unsigned int lit;
    unsigned int show_n;
    uint8_t active;
} strip_slot_t;

static strip_slot_t s_slots[MAX_STRIPS];
static unsigned int s_max_leds;

static int claim_sm(PIO *pio_out, uint *sm_out) {
    int sm = pio_claim_unused_sm(pio0, false);
    if (sm >= 0) {
        *pio_out = pio0;
        *sm_out = (uint)sm;
        return 0;
    }
    sm = pio_claim_unused_sm(pio1, false);
    if (sm >= 0) {
        *pio_out = pio1;
        *sm_out = (uint)sm;
        return 0;
    }
    return -1;
}

extern "C" void ws2812_init(unsigned int max_leds) {
    s_max_leds = max_leds;
    for (unsigned int i = 0; i < MAX_STRIPS; i++) {
        s_slots[i].strip = nullptr;
        s_slots[i].pin = 0;
        s_slots[i].lit = 0;
        s_slots[i].show_n = 0;
        s_slots[i].active = 0;
    }
}

extern "C" void ws2812_release_strip(unsigned int strip) {
    if (strip >= MAX_STRIPS) {
        return;
    }
    strip_slot_t *slot = &s_slots[strip];
    if (!slot->active) {
        return;
    }
    pio_sm_set_enabled(slot->pio, slot->sm, false);
    gpio_set_function(slot->pin, GPIO_FUNC_SIO);
    gpio_set_dir(slot->pin, GPIO_IN);
    delete slot->strip;
    slot->strip = nullptr;
    pio_sm_unclaim(slot->pio, slot->sm);
    slot->pin = 0;
    slot->lit = 0;
    slot->show_n = 0;
    slot->active = 0;
}

extern "C" int ws2812_set_pin(unsigned int strip, unsigned int pin) {
    if (strip >= MAX_STRIPS || pin > 29) {
        return -1;
    }
    strip_slot_t *slot = &s_slots[strip];
    if (slot->active && slot->pin == pin) {
        return 0;
    }
    /* Steal pin from any other strip that already owns it. */
    for (unsigned int s = 0; s < MAX_STRIPS; s++) {
        if (s != strip && s_slots[s].active && s_slots[s].pin == pin) {
            ws2812_release_strip(s);
        }
    }
    if (slot->active) {
        /* Tear down and rebuild on the same PIO SM — more reliable than setPin alone. */
        PIO pio = slot->pio;
        uint sm = slot->sm;
        unsigned int old_pin = slot->pin;
        pio_sm_set_enabled(pio, sm, false);
        gpio_set_function(old_pin, GPIO_FUNC_SIO);
        gpio_set_dir(old_pin, GPIO_IN);
        delete slot->strip;
        slot->strip = new WS2812(pin, s_max_leds, pio, sm, WS2812::FORMAT_GRB);
        slot->pin = pin;
        slot->lit = 0;
        slot->show_n = 0;
        return 0;
    }
    PIO pio;
    uint sm;
    if (claim_sm(&pio, &sm) != 0) {
        return -1;
    }
    slot->strip = new WS2812(pin, s_max_leds, pio, sm, WS2812::FORMAT_GRB);
    slot->pio = pio;
    slot->sm = sm;
    slot->pin = pin;
    slot->lit = 0;
    slot->show_n = 0;
    slot->active = 1;
    return 0;
}

extern "C" unsigned int ws2812_get_pin(unsigned int strip) {
    if (strip >= MAX_STRIPS || !s_slots[strip].active) {
        return 0;
    }
    return s_slots[strip].pin;
}

extern "C" int ws2812_strip_active(unsigned int strip) {
    return strip < MAX_STRIPS && s_slots[strip].active;
}

extern "C" void ws2812_set_rgb(unsigned int strip, const uint8_t *rgb, unsigned int count) {
    if (strip >= MAX_STRIPS || !s_slots[strip].active || rgb == nullptr) {
        return;
    }
    strip_slot_t *slot = &s_slots[strip];
    if (count > s_max_leds) {
        count = s_max_leds;
    }
    unsigned int old = slot->lit;
    unsigned int i = 0;
    for (; i < count; i++) {
        slot->strip->setPixelColor(i, rgb[i * 3 + 0], rgb[i * 3 + 1], rgb[i * 3 + 2]);
    }
    for (; i < old; i++) {
        slot->strip->setPixelColor(i, 0, 0, 0);
    }
    slot->lit = count;
    slot->show_n = count > old ? count : old;
}

extern "C" void ws2812_fill(unsigned int strip, uint8_t r, uint8_t g, uint8_t b) {
    if (strip >= MAX_STRIPS || !s_slots[strip].active) {
        return;
    }
    strip_slot_t *slot = &s_slots[strip];
    slot->strip->fill(WS2812::RGB(r, g, b));
    slot->lit = s_max_leds;
    slot->show_n = s_max_leds;
}

extern "C" void ws2812_show(unsigned int strip) {
    if (strip >= MAX_STRIPS || !s_slots[strip].active) {
        return;
    }
    strip_slot_t *slot = &s_slots[strip];
    unsigned int n = slot->show_n;
    if (n == 0) {
        n = 1;
    }
    slot->strip->show(n);
    sleep_us(300);
    slot->show_n = slot->lit;
}

extern "C" void ws2812_show_all(void) {
    for (unsigned int i = 0; i < MAX_STRIPS; i++) {
        if (s_slots[i].active) {
            ws2812_show(i);
        }
    }
}