neiam /lambent-flex-c6

use anyhow::Result;
use core::time::Duration;
use esp_idf_svc::hal::{
    gpio::OutputPin,
    rmt::{
        config::TransmitConfig, PinState, Pulse, RmtChannel, TxRmtDriver, VariableLengthSignal,
    },
};

pub use rgb::RGB8;

/// Bytes per pixel on the wire.
const BPP: usize = 3;
/// RMT pulse pairs per pixel (24 bits, one high+low pair each).
const PULSES_PER_PIXEL: usize = 24;

/// Ceiling on pixels per frame, sized to the UDP handler's 1024-byte buffer.
/// Anything longer than the buffer is truncated by `recv_from` before it ever
/// reaches here.
const MAX_PIXELS: usize = 1024 / BPP;

pub struct WS2812RMT<'a> {
    tx_rtm_driver: TxRmtDriver<'a>,
    /// Cached at construction: `counter_clock()` does not change, and rebuilding
    /// these four per frame was pure overhead.
    t0: (Pulse, Pulse),
    t1: (Pulse, Pulse),
    /// Reused across frames so a 341-pixel write does not reallocate every packet.
    signal: VariableLengthSignal,
}

impl<'d> WS2812RMT<'d> {
    // Rust ESP Board gpio2, ESP32-C6-DevKitC-1 gpio8
    pub fn new(
        // esp-idf-hal 0.46 removed the `Peripheral` indirection; drivers now take
        // the peripheral by value directly.
        led: impl OutputPin + 'd,
        channel: impl RmtChannel + 'd,
    ) -> Result<Self> {
        let config = TransmitConfig::new().clock_divider(2);
        let tx = TxRmtDriver::new(channel, led, &config)?;

        let ticks_hz = tx.counter_clock()?;
        let t0 = (
            Pulse::new_with_duration(ticks_hz, PinState::High, &ns(350))?,
            Pulse::new_with_duration(ticks_hz, PinState::Low, &ns(800))?,
        );
        let t1 = (
            Pulse::new_with_duration(ticks_hz, PinState::High, &ns(700))?,
            Pulse::new_with_duration(ticks_hz, PinState::Low, &ns(600))?,
        );

        Ok(Self {
            tx_rtm_driver: tx,
            t0,
            t1,
            signal: VariableLengthSignal::with_capacity(MAX_PIXELS * PULSES_PER_PIXEL),
        })
    }

    /// Drive a whole strand from a packed `R,G,B,R,G,B,...` buffer.
    ///
    /// One pixel per 3 bytes, in order. A trailing partial triple is ignored —
    /// there is no sane colour to invent for it, and dropping it keeps a
    /// truncated packet from shifting every later pixel.
    ///
    /// Pixels beyond the end of the physical strand are simply not latched by
    /// any LED, so overshooting is harmless; undershooting leaves the remaining
    /// LEDs holding their previous colour, because WS2812 only updates pixels
    /// that actually receive data.
    pub fn set_pixels(&mut self, data: &[u8]) -> Result<()> {
        let (t0h, t0l) = self.t0;
        let (t1h, t1l) = self.t1;

        self.signal.clear();
        for px in data.chunks_exact(BPP).take(MAX_PIXELS) {
            // Wire order is GRB, while the packet is RGB.
            let color: u32 =
                ((px[1] as u32) << 16) | ((px[0] as u32) << 8) | px[2] as u32;

            for i in (0..PULSES_PER_PIXEL).rev() {
                let (high, low) = if color & (1 << i) != 0 {
                    (t1h, t1l)
                } else {
                    (t0h, t0l)
                };
                self.signal.push([&high, &low])?;
            }
        }

        // Nothing to send: an empty signal is not a valid RMT transmission, and
        // there is no state to update anyway.
        if data.len() < BPP {
            return Ok(());
        }

        self.tx_rtm_driver.start_blocking(&self.signal)?;
        Ok(())
    }

    /// Single-pixel convenience wrapper, kept for callers that only have one
    /// colour to send.
    pub fn set_pixel(&mut self, rgb: RGB8) -> Result<()> {
        self.set_pixels(&[rgb.r, rgb.g, rgb.b])
    }
}

fn ns(nanos: u64) -> Duration {
    Duration::from_nanos(nanos)
}