import array, time
from machine import Pin
import rp2
from time import sleep
import dma

@rp2.asm_pio(sideset_init=rp2.PIO.OUT_LOW, out_shiftdir=rp2.PIO.SHIFT_LEFT, autopull=True, pull_thresh=8)
def ws2812():
    T1 = 2
    T2 = 5
    T3 = 3
    wrap_target()
    label("bitloop")
    out(x, 1)               .side(0)    [T3 - 1]
    jmp(not_x, "do_zero")   .side(1)    [T1 - 1]
    jmp("bitloop")          .side(1)    [T2 - 1]
    label("do_zero")
    nop()                   .side(0)    [T2 - 1]
    wrap()

class WS2812B:
    BLACK = (0, 0, 0)
    RED = (255, 0, 0)
    YELLOW = (255, 150, 0)
    GREEN = (0, 255, 0)
    CYAN = (0, 255, 255)
    BLUE = (0, 0, 255)
    PURPLE = (180, 0, 255)
    WHITE = (255, 255, 255)
    COLORS = (BLACK, RED, YELLOW, GREEN, CYAN, BLUE, PURPLE, WHITE)

    def __init__(self, num_leds, pin, state_machine, brightness=0.1):
        self.sm = rp2.StateMachine(state_machine, ws2812, freq=8_000_000, sideset_base=Pin(pin))
        self.sm.active(1)
        self.ar = bytearray(num_leds * 3)
        self.num_leds = num_leds
        self.brightness = brightness
        self.pio_dma = dma.PIO_DMA_Transfer(state_machine, state_machine, 8, num_leds * 3)
        # Pre-built bytearrays, one per color (GRB, brightness applied) for fast show
        self.color_buffers = [self._color_to_buffer(c) for c in self.COLORS]

    def _color_to_buffer(self, color):
        """One bytearray of length num_leds*3, all pixels same color (GRB)."""
        r, g, b = color[0], color[1], color[2]
        g = int(g * self.brightness) & 0xFF
        r = int(r * self.brightness) & 0xFF
        b = int(b * self.brightness) & 0xFF
        return bytearray([g, r, b] * self.num_leds)

    def show(self, buffer=None, offset=0):
        """Push buffer to PIO via DMA. If buffer is None, use self.ar (offset must be 0).
        With offset, DMA reads directly from buffer[offset:offset+num_leds*3]; no copy."""
        buf = buffer if buffer is not None else self.ar
        self.pio_dma.start_transfer(buf, offset)

    def show_color(self, color_index):
        """Show pre-built buffer for COLORS[color_index]. No fill() needed."""
        self.show(self.color_buffers[color_index])

    def set(self, i, color):
        self.ar[i*3] = int(color[1]*self.brightness)
        self.ar[i*3+1] = int(color[0]*self.brightness)
        self.ar[i*3+2] = int(color[2]*self.brightness)

    def fill(self, color):
        for i in range(self.num_leds):
            self.set(i, color)

    def busy(self):
        return self.pio_dma.busy()

if __name__ == "__main__":
    num_leds, pin, sm, brightness = 10, 2, 0, 1
    ws0 = WS2812B(num_leds, pin, sm, brightness)
    while True:
        for color in ws0.COLORS:
            ws0.fill(color)
            ws0.show()
            time.sleep(1)