import array
import time
from machine import Pin
from rp2 import StateMachine, DMA, PIO, asm_pio
from uctypes import addressof

################## DMX OUT ##################
@asm_pio(
    set_init=PIO.OUT_LOW,
    sideset_init=PIO.OUT_LOW,
    out_init=PIO.OUT_LOW,
    out_shiftdir=PIO.SHIFT_RIGHT  
)
def dmx_tx():
    """
    First word = packet Length
    Second word = START_CODE
    512 words of channel data
    """
    wrap_target()
    # Get Packet Length (typically 513)    
    pull()                                      # Get a 32-bit word from the TX FIFO
    mov(x, osr)                                 # x now holds the full loop count (e.g. 513)
    
    # Break                                
    set(y, 24)             .side(0)             # BREAK duration (24 iterations (0–22) × 4 cycles = 100 µs)    
    label("breakloop")                                   
    jmp(y_dec, "breakloop")           [7]       # 4 cycles (1 instruction + 3 delay) 

    # MAB (Mark After Break)
    nop()                  .side(1)   [7]       # 12 cycles (1 instruction + 11 delay)

    # Send Packet
    label("byteloop")
    pull()                 .side(1)   [7]       # Get a 32-bit word from the TX FIFO + Start bit
    
    # Send byte (1 start bit, 8 data bits, 2 stop bits)
    set(y, 7)              .side(0)   [3]       # 8 bits to send
    label("bitloop")
    out(pins, 1)                      [2]       # 4 µs = 4 cycles (1 instruction + 3 delay)
    jmp(y_dec, "bitloop")                       # Loop 8 times for each bit: 8 * 4 = 32µs
    
    jmp(x_dec, "byteloop")                      # Send next byte
    wrap()
    
class DMX:
    """
    DMX using PIO and chained DMA on RP2040.
    
    Args:
        tx_pin (Pin): Pin object for DMX transmit line.
        sm_id (int): State Machine ID (0-7).
        buffer (array.array): Array of 32-bit words containing DMX data, first word is packet length (IE 513).
        ctrl_pin (Pin, optional): Pin to control RS485 direction.
    """    
    def __init__(self, tx_pin: Pin, sm_id: int, ctrl_pin: Pin = None):
        self.out_pin = Pin(tx_pin, Pin.OUT)       # Configure pins to be output
        self.sm_id = sm_id 
        self.buffer = array.array('I', [513] + [0] * 513)         
    
        self.start_sm_dma()   # State Machine - Setup & Start 
        
    def start_sm_dma(self):
        # Configure and start State Machine with settings
        self.sm = StateMachine(
            self.sm_id,
            dmx_tx,
            freq=1_000_000,
            out_base=self.out_pin,
            set_base=self.out_pin,
            sideset_base=self.out_pin
        )        
        self.sm.active(1)  # Start SM

        # Configure and Start chained DMA channels with settings       
        self.dma = [DMA(), DMA()]
        
        # This channel holds the DMX data and if fed to PIO FIFO    
        self.dma[0].config( 
            read = self.buffer, 
            write = self.sm,
            count = len(self.buffer),  # 514 words
            ctrl = self.dma[0].pack_ctrl(
                size = 2,  # 0 = byte (8 bit), 1 = half word (16 bit), 2 = word (32 bit)
                inc_read=True,
                inc_write=False,
                treq_sel=self.sm_id,
                chain_to=self.dma[1].channel
            ),
        ) 
    
        # This channel just retriggers the previous channel to create a loop.
        # https://github.com/orgs/micropython/discussions/14420
        self.addresses_of_data = array.array("I", [addressof(self.buffer)])
        
        self.dma[1].config(
            read = self.addresses_of_data,
            write = self.dma[0].registers[15:16],  # 15 is CHx_AL3_READ_ADDR_TRIG
            count = 1,
            ctrl = self.dma[1].pack_ctrl(   # read/write only from/to one address
                size = 2,  # 0 = byte (8 bit), 1 = half word (16 bit), 2 = word (32 bit)               
                inc_read=False,
                inc_write=False,
                chain_to=self.dma[1].channel  # loops back to self - disables chaining
            ),
        )
        
        self.dma[0].active(1)  # Start DM        
        
    def close(self):
        # Tear down DMA channels
        for d in (self.dma[1], self.dma[0]):  # Reverse order to handle chaining
            d.active(0)
            d.close()       # Release DMA resources
        
        # Stop PIO State Machine
        self.sm.active(0)        


if __name__ == "__main__":               
    dmx_instance = DMX(tx_pin=0, sm_id=0)
    packet_len = 513
    
    try:
        while True:
            dmx_instance.buffer[0] = packet_len # Set packet length into buffer                        
            dmx_instance.buffer[1:] = array.array('I', [i & 0xFF for i in range(packet_len)])  # Set dummy data into buffer
            
            #print(f"DMX: {list(dmx_instance.buffer[:8])}")            
            time.sleep_ms(40)
            
    except KeyboardInterrupt:
        print("Interrupt detected. Cleaning up...")
    finally:
        dmx_instance.close()
        print("Program terminated.")