"""
Pico 2 LED Controller for Vixy's Head

Dual WS2812 strip control via USB serial.
- Strip 0 (Jaw):  GP0, 14 LEDs (index 0 damaged)
- Strip 1 (Mood): GP1, 56 LEDs

Protocol:
  "0 index r g b"  - Set jaw LED
  "1 index r g b"  - Set mood LED  
  "0 -1"           - Trigger jaw output
  "1 -1"           - Trigger mood output
  "0 clear"        - Clear jaw
  "1 clear"        - Clear mood
  "1 fill r g b"   - Fill mood strip
"""

import sys
import array
import time
from machine import Pin
import rp2

# Configuration
JAW_PIN = 0
MOOD_PIN = 1
JAW_LEDS = 14
MOOD_LEDS = 56

# WS2812 PIO program - 800kHz timing
@rp2.asm_pio(sideset_init=rp2.PIO.OUT_LOW, out_shiftdir=rp2.PIO.SHIFT_LEFT, autopull=True, pull_thresh=24)
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 WS2812Strip:
    def __init__(self, pin, num_leds, sm_id):
        self.num_leds = num_leds
        self.sm = rp2.StateMachine(sm_id, ws2812, freq=8_000_000, sideset_base=Pin(pin))
        self.sm.active(1)
        self.ar = array.array("I", [0] * num_leds)
    
    def set_led(self, index, r, g, b):
        if 0 <= index < self.num_leds:
            self.ar[index] = (g << 16) | (r << 8) | b
    
    def fill(self, r, g, b):
        color = (g << 16) | (r << 8) | b
        for i in range(self.num_leds):
            self.ar[i] = color
    
    def clear(self):
        for i in range(self.num_leds):
            self.ar[i] = 0
    
    def show(self):
        for c in self.ar:
            self.sm.put(c, 8)
        time.sleep_us(60)


# Initialize strips
jaw = WS2812Strip(JAW_PIN, JAW_LEDS, 0)
mood = WS2812Strip(MOOD_PIN, MOOD_LEDS, 1)
strips = [jaw, mood]

# Clear both on startup
jaw.clear()
jaw.show()
mood.clear()
mood.show()

print("Pico LED Controller ready")
print(f"Jaw: GP{JAW_PIN}, {JAW_LEDS} LEDs")
print(f"Mood: GP{MOOD_PIN}, {MOOD_LEDS} LEDs")


def process_command(line):
    """Parse and execute a command."""
    parts = line.strip().split()
    if len(parts) < 2:
        return
    
    try:
        strip_id = int(parts[0])
        if strip_id not in (0, 1):
            return
        
        strip = strips[strip_id]
        cmd = parts[1]
        
        if cmd == "clear":
            strip.clear()
            strip.show()
            print("OK")
        
        elif cmd == "-1":
            strip.show()
            print("OK")
        
        elif cmd == "fill" and len(parts) >= 5:
            r, g, b = int(parts[2]), int(parts[3]), int(parts[4])
            strip.fill(r, g, b)
            strip.show()
            print("OK")
        
        elif len(parts) >= 5:
            index = int(cmd)
            r, g, b = int(parts[2]), int(parts[3]), int(parts[4])
            strip.set_led(index, r, g, b)
            print("OK")
    
    except (ValueError, IndexError) as e:
        print(f"ERR: {e}")


# Main loop - blocking readline is fine for this use case
while True:
    try:
        line = sys.stdin.readline()
        if line:
            process_command(line)
    except Exception as e:
        print(f"ERR: {e}")