piopulse 0.1.3

A terminal user interface (TUI) factory flashing tool designed for high-concurrency ESP32 chip flashing on production lines.
use std::sync::atomic::AtomicU8;

pub static ACTIVE_VOFA_MODE: AtomicU8 = AtomicU8::new(0); // 0 = FireWater, 1 = JustFloat, 2 = IndexFloat

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum VofaMode {
    FireWater = 0,
    JustFloat = 1,
    IndexFloat = 2,
}

impl VofaMode {
    pub fn to_u8(self) -> u8 {
        self as u8
    }

    pub fn from_u8(val: u8) -> Self {
        match val {
            0 => VofaMode::FireWater,
            1 => VofaMode::JustFloat,
            _ => VofaMode::IndexFloat,
        }
    }
}

#[derive(Debug, Clone)]
pub struct VofaImage {
    pub id: usize,
    pub width: usize,
    pub height: usize,
    pub format: u8,
    pub data: Vec<u8>,
}

pub struct VofaParser {
    buffer: Vec<u8>,
    mode: VofaMode,
    index_float_buffer: Vec<f32>,
    pub latest_image: Option<VofaImage>,
}

impl VofaParser {
    pub fn new(mode: VofaMode) -> Self {
        Self {
            buffer: Vec::new(),
            mode,
            index_float_buffer: Vec::new(),
            latest_image: None,
        }
    }

    pub fn take_latest_image(&mut self) -> Option<VofaImage> {
        self.latest_image.take()
    }

    pub fn set_mode(&mut self, mode: VofaMode) {
        if self.mode != mode {
            self.mode = mode;
            self.buffer.clear();
            self.index_float_buffer.clear();
        }
    }

    pub fn feed(&mut self, data: &[u8]) -> Vec<Vec<f32>> {
        self.buffer.extend_from_slice(data);
        let mut frames = Vec::new();

        loop {
            match self.mode {
                VofaMode::FireWater => {
                    // FireWater is CSV-like: val1,val2,val3...\n (or \r\n)
                    if let Some(pos) = self.buffer.iter().position(|&b| b == b'\n') {
                        let is_image_frame =
                            if let Ok(line_str) = std::str::from_utf8(&self.buffer[..pos]) {
                                let trimmed = line_str.trim();
                                if let Some(colon_pos) = trimmed.find(':') {
                                    let name = trimmed[..colon_pos].trim();
                                    if name == "image" {
                                        let datas: Vec<usize> = trimmed[colon_pos + 1..]
                                            .split(',')
                                            .filter_map(|s| s.trim().parse::<usize>().ok())
                                            .collect();
                                        if datas.len() == 5 {
                                            Some((datas[0], datas[1], datas[2], datas[3], datas[4]))
                                        } else {
                                            None
                                        }
                                    } else {
                                        None
                                    }
                                } else {
                                    None
                                }
                            } else {
                                None
                            };

                        if let Some((
                            image_id,
                            image_size,
                            image_width,
                            image_height,
                            image_format,
                        )) = is_image_frame
                        {
                            if self.buffer.len() < pos + 1 + image_size {
                                break; // Wait for full image data
                            }
                            self.buffer.drain(..=pos);
                            let image_data = self.buffer.drain(..image_size).collect::<Vec<u8>>();
                            self.latest_image = Some(VofaImage {
                                id: image_id,
                                width: image_width,
                                height: image_height,
                                format: image_format as u8,
                                data: image_data,
                            });
                        } else {
                            let line_bytes = self.buffer.drain(..=pos).collect::<Vec<u8>>();
                            if let Ok(line_str) = std::str::from_utf8(&line_bytes) {
                                let trimmed = line_str.trim();
                                let data_part = if let Some(colon_pos) = trimmed.find(':') {
                                    &trimmed[colon_pos + 1..]
                                } else {
                                    trimmed
                                };
                                let parsed: Vec<f32> = data_part
                                    .split(',')
                                    .filter_map(|s| s.trim().parse::<f32>().ok())
                                    .collect();
                                if !parsed.is_empty() {
                                    frames.push(parsed);
                                }
                            }
                        }
                    } else {
                        break;
                    }
                }
                VofaMode::JustFloat => {
                    // JustFloat is binary floats ending with NaN: [0x00, 0x00, 0x80, 0x7F]
                    let tail = [0x00, 0x00, 0x80, 0x7F];
                    if let Some(pos) = self.buffer.windows(4).position(|w| w == tail) {
                        let is_image = if pos == 20 {
                            if self.buffer.len() >= pos + 8 {
                                self.buffer[pos + 4..pos + 8] == tail
                            } else {
                                // Heuristics on metadata to see if we should wait for the second NaN
                                let size = i32::from_le_bytes(
                                    self.buffer[4..8].try_into().unwrap_or([0; 4]),
                                );
                                let width = i32::from_le_bytes(
                                    self.buffer[8..12].try_into().unwrap_or([0; 4]),
                                );
                                let height = i32::from_le_bytes(
                                    self.buffer[12..16].try_into().unwrap_or([0; 4]),
                                );
                                let format = i32::from_le_bytes(
                                    self.buffer[16..20].try_into().unwrap_or([0; 4]),
                                );

                                let looks_like_image_header = size > 0
                                    && (width > 0 || width == -1)
                                    && (height > 0 || height == -1)
                                    && (format >= 0 && format <= 34);

                                if looks_like_image_header {
                                    break; // wait for more data
                                }
                                false
                            }
                        } else {
                            false
                        };

                        if is_image {
                            let image_id =
                                i32::from_le_bytes(self.buffer[0..4].try_into().unwrap_or([0; 4]))
                                    as usize;
                            let image_size =
                                i32::from_le_bytes(self.buffer[4..8].try_into().unwrap_or([0; 4]))
                                    as usize;
                            let image_width =
                                i32::from_le_bytes(self.buffer[8..12].try_into().unwrap_or([0; 4]))
                                    as usize;
                            let image_height = i32::from_le_bytes(
                                self.buffer[12..16].try_into().unwrap_or([0; 4]),
                            ) as usize;
                            let image_format = i32::from_le_bytes(
                                self.buffer[16..20].try_into().unwrap_or([0; 4]),
                            ) as u8;

                            if self.buffer.len() < pos + 8 + image_size {
                                break; // Wait for full image payload
                            }

                            self.buffer.drain(..pos + 8);
                            let image_data = self.buffer.drain(..image_size).collect::<Vec<u8>>();
                            self.latest_image = Some(VofaImage {
                                id: image_id,
                                width: image_width,
                                height: image_height,
                                format: image_format,
                                data: image_data,
                            });
                        } else {
                            let frame_bytes = self.buffer.drain(..pos + 4).collect::<Vec<u8>>();
                            let val_bytes = &frame_bytes[..frame_bytes.len() - 4];
                            let mut vals = Vec::new();
                            for chunk in val_bytes.chunks_exact(4) {
                                let val = f32::from_le_bytes(chunk.try_into().unwrap());
                                vals.push(val);
                            }
                            if !vals.is_empty() {
                                frames.push(vals);
                            }
                        }
                    } else {
                        break;
                    }
                }
                VofaMode::IndexFloat => {
                    // IndexFloat is binary floats ending with NaN: [0x00, 0x00, 0x80, 0x7F]
                    let tail = [0x00, 0x00, 0x80, 0x7F];
                    if let Some(pos) = self.buffer.windows(4).position(|w| w == tail) {
                        let frame_bytes = self.buffer.drain(..pos + 4).collect::<Vec<u8>>();
                        if frame_bytes.len() >= 8 && frame_bytes.len() % 4 == 0 {
                            let val_bytes = &frame_bytes[..frame_bytes.len() - 4];

                            let start_index_val =
                                f32::from_le_bytes(val_bytes[0..4].try_into().unwrap());
                            let start_index = start_index_val as usize;

                            let data_bytes = &val_bytes[4..];
                            let data_count = data_bytes.len() / 4;

                            if self.index_float_buffer.len() < start_index + data_count {
                                self.index_float_buffer
                                    .resize(start_index + data_count, 0.0);
                            }

                            for (idx, chunk) in data_bytes.chunks_exact(4).enumerate() {
                                let val = f32::from_le_bytes(chunk.try_into().unwrap());
                                self.index_float_buffer[start_index + idx] = val;
                            }

                            frames.push(self.index_float_buffer.clone());
                        }
                    } else {
                        break;
                    }
                }
            }
        }
        frames
    }

    #[allow(dead_code)]
    pub fn clear(&mut self) {
        self.buffer.clear();
        self.index_float_buffer.clear();
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_firewater_parser() {
        let mut parser = VofaParser::new(VofaMode::FireWater);

        // Test normal CSV
        let res = parser.feed(b"1.2,3.4\n");
        assert_eq!(res, vec![vec![1.2, 3.4]]);

        // Test CSV with label prefix
        let res = parser.feed(b"d:5.6,7.8\n");
        assert_eq!(res, vec![vec![5.6, 7.8]]);

        // Test CSV with spaces and CRLF
        let res = parser.feed(b"  my_label :  10.0 , 20.0 \r\n");
        assert_eq!(res, vec![vec![10.0, 20.0]]);
    }

    #[test]
    fn test_firewater_image_parser() {
        let mut parser = VofaParser::new(VofaMode::FireWater);

        // Feed partial image data header and verify no image yet
        let res = parser.feed(b"image:0,6,2,3,0\nabcd");
        assert!(res.is_empty());
        assert!(parser.take_latest_image().is_none());

        // Feed the rest of the image data
        let res = parser.feed(b"ef");
        assert!(res.is_empty());

        let img = parser.take_latest_image().expect("Should parse image");
        assert_eq!(img.id, 0);
        assert_eq!(img.width, 2);
        assert_eq!(img.height, 3);
        assert_eq!(img.format, 0);
        assert_eq!(img.data, b"abcdef");
    }

    #[test]
    fn test_justfloat_image_parser() {
        let mut parser = VofaParser::new(VofaMode::JustFloat);

        // Pre-frame variables
        let img_id: i32 = 1;
        let img_size: i32 = 6;
        let img_width: i32 = 2;
        let img_height: i32 = 3;
        let img_format: i32 = 24; // Grayscale8

        let mut buf = Vec::new();
        buf.extend_from_slice(&img_id.to_le_bytes());
        buf.extend_from_slice(&img_size.to_le_bytes());
        buf.extend_from_slice(&img_width.to_le_bytes());
        buf.extend_from_slice(&img_height.to_le_bytes());
        buf.extend_from_slice(&img_format.to_le_bytes());
        buf.extend_from_slice(&[0x00, 0x00, 0x80, 0x7F]); // First NaN
        buf.extend_from_slice(&[0x00, 0x00, 0x80, 0x7F]); // Second NaN

        // Feed partial pre-frame data and verify no image yet
        let res = parser.feed(&buf[..buf.len() - 4]); // omit last 4 bytes of NaN
        assert!(res.is_empty());
        assert!(parser.take_latest_image().is_none());

        // Feed rest of preframe and partial payload (4 bytes of b"abcdef")
        let mut part2 = Vec::new();
        part2.extend_from_slice(&[0x00, 0x00, 0x80, 0x7F]); // complete header
        part2.extend_from_slice(b"abcd");
        let res = parser.feed(&part2);
        assert!(res.is_empty());
        assert!(parser.take_latest_image().is_none());

        // Feed final 2 bytes of payload
        let res = parser.feed(b"ef");
        assert!(res.is_empty());

        let img = parser.take_latest_image().expect("Should parse image");
        assert_eq!(img.id, 1);
        assert_eq!(img.width, 2);
        assert_eq!(img.height, 3);
        assert_eq!(img.format, 24);
        assert_eq!(img.data, b"abcdef");
    }
}