trx-rs/src/decoders/trx-sstv/src/image.rs

// SPDX-FileCopyrightText: 2026 Stan Grams <sjg@haxx.space>
//
// SPDX-License-Identifier: GPL-2.0-or-later

//! Assembling decoded lines into an image, and getting it out of the process.
//!
//! Rows arrive one at a time and the picture is worth looking at before it is
//! finished, so the assembler holds a full-size RGB canvas from the start and
//! fills it in. An unfinished frame is grey below the last decoded row rather
//! than black, which reads as "not here yet" instead of "received as black".

use std::path::{Path, PathBuf};

use base64::Engine;

/// Value the canvas starts at: mid-grey, for rows not yet received.
const UNWRITTEN: u8 = 96;

pub struct ImageCanvas {
    width: usize,
    height: usize,
    rgb: Vec<u8>,
    /// Highest row index written, plus one.
    filled_rows: usize,
}

impl ImageCanvas {
    pub fn new(width: usize, height: usize) -> Self {
        Self {
            width,
            height,
            rgb: vec![UNWRITTEN; width * height * 3],
            filled_rows: 0,
        }
    }

    pub fn width(&self) -> usize {
        self.width
    }

    pub fn height(&self) -> usize {
        self.height
    }

    /// Rows written so far.
    pub fn filled_rows(&self) -> usize {
        self.filled_rows
    }

    /// Write one row of RGB triples. Rows past the bottom of the image are
    /// dropped: a transmission that runs long is not a reason to grow.
    pub fn put_row(&mut self, y: usize, row: &[u8]) {
        if y >= self.height {
            return;
        }
        let at = y * self.width * 3;
        let take = row.len().min(self.width * 3);
        self.rgb[at..at + take].copy_from_slice(&row[..take]);
        self.filled_rows = self.filled_rows.max(y + 1);
    }

    pub fn row(&self, y: usize) -> Option<&[u8]> {
        if y >= self.height {
            return None;
        }
        let at = y * self.width * 3;
        Some(&self.rgb[at..at + self.width * 3])
    }

    pub fn rgb(&self) -> &[u8] {
        &self.rgb
    }

    /// Encode the canvas as a PNG.
    pub fn to_png(&self) -> Result<Vec<u8>, String> {
        let mut out = Vec::new();
        {
            let mut encoder = png::Encoder::new(&mut out, self.width as u32, self.height as u32);
            encoder.set_color(png::ColorType::Rgb);
            encoder.set_depth(png::BitDepth::Eight);
            let mut writer = encoder
                .write_header()
                .map_err(|e| format!("PNG header: {e}"))?;
            writer
                .write_image_data(&self.rgb)
                .map_err(|e| format!("PNG data: {e}"))?;
        }
        Ok(out)
    }

    /// The PNG, base64-encoded for the journey to a browser.
    pub fn to_png_base64(&self) -> Result<String, String> {
        Ok(base64::engine::general_purpose::STANDARD.encode(self.to_png()?))
    }

    /// Write the PNG into `dir`, named for when and where it was received.
    pub fn save_png(
        &self,
        dir: &Path,
        freq_hz: u64,
        mode_name: &str,
        stamp: &str,
    ) -> Result<PathBuf, String> {
        std::fs::create_dir_all(dir).map_err(|e| format!("create {}: {e}", dir.display()))?;
        let slug: String = mode_name
            .chars()
            .map(|c| if c.is_ascii_alphanumeric() { c } else { '-' })
            .collect();
        let path = dir.join(format!("SSTV_{stamp}_{freq_hz}_{slug}.png"));
        std::fs::write(&path, self.to_png()?)
            .map_err(|e| format!("write {}: {e}", path.display()))?;
        Ok(path)
    }
}

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

    #[test]
    fn rows_land_where_they_are_put_and_the_rest_stays_unwritten() {
        let mut canvas = ImageCanvas::new(4, 3);
        canvas.put_row(1, &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]);
        assert_eq!(canvas.row(1).unwrap()[0..3], [1, 2, 3]);
        assert_eq!(canvas.row(0).unwrap()[0], UNWRITTEN);
        assert_eq!(canvas.filled_rows(), 2);
    }

    #[test]
    fn a_row_past_the_bottom_is_dropped_rather_than_growing_the_image() {
        let mut canvas = ImageCanvas::new(2, 2);
        canvas.put_row(9, &[1, 2, 3, 4, 5, 6]);
        assert_eq!(canvas.filled_rows(), 0);
        assert_eq!(canvas.rgb().len(), 2 * 2 * 3);
    }

    #[test]
    fn encodes_a_png_a_decoder_can_read_back() {
        let mut canvas = ImageCanvas::new(2, 2);
        canvas.put_row(0, &[255, 0, 0, 0, 255, 0]);
        let png_bytes = canvas.to_png().expect("png");
        let decoder = png::Decoder::new(png_bytes.as_slice());
        let mut reader = decoder.read_info().expect("png info");
        let mut buf = vec![0; reader.output_buffer_size()];
        let info = reader.next_frame(&mut buf).expect("png frame");
        assert_eq!((info.width, info.height), (2, 2));
        assert_eq!(&buf[0..6], &[255, 0, 0, 0, 255, 0]);
        assert!(!canvas.to_png_base64().expect("base64").is_empty());
    }
}