From 85a6cd516341c36775b02207117277566639fd90 Mon Sep 17 00:00:00 2001
From: Joakim Hulthe <joakim@hulthe.net>
Date: Mon, 14 Jul 2025 17:31:06 +0200
Subject: [PATCH] Add --stamp flag

---
 src/main.rs | 110 ++++++++++++++++++++++++++++++++++++++++++----------
 1 file changed, 89 insertions(+), 21 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index b288711..d8be526 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,23 +1,31 @@
-//! Image format is a buffer of `SCREEN_W * SCREEN_H + 1` bytes.
-//! The 4 least-significant bits of each byte is one pixel. The 4 most-significant bits are unused.
-//! Images are stored in landscape format, i.e. the first byte is the top-left pixel, and the
-//! second-to-last byte is the bottom-right pixel. The final byte is a null-terminator, and must be 0.
 use clap::Parser;
-use eyre::{Context, bail, eyre};
-use image::ImageReader;
+use eyre::{Context, ContextCompat, bail, eyre};
+use image::{DynamicImage, GrayImage, ImageReader, imageops::overlay};
 use nix::ioctl_write_ptr;
-use std::{fs::File, os::fd::AsRawFd, path::PathBuf};
+use std::{
+    fs::File,
+    os::fd::AsRawFd,
+    path::{Path, PathBuf},
+};
 
 #[derive(Parser)]
 struct Opt {
     /// Path to the image-blob.
     pub path: PathBuf,
+
+    /// Stamp an image onto the final image. "<x>:<y>:<path>"
+    #[clap(long)]
+    pub stamp: Vec<String>,
 }
 
 const SCREEN_W: usize = 1872;
 const SCREEN_H: usize = 1404;
 
-const FILE_SIZE: usize = SCREEN_W * SCREEN_H + 1;
+/// Image format is a buffer of `SCREEN_W * SCREEN_H + 1` bytes.
+/// The 4 least-significant bits of each byte is one pixel. The 4 most-significant bits are unused.
+/// Images are stored in landscape format, i.e. the first byte is the top-left pixel, and the
+/// second-to-last byte is the bottom-right pixel. The final byte is a null-terminator, and must be 0.
+const IOCTL_BUFFER_SIZE: usize = SCREEN_W * SCREEN_H + 1;
 
 // const SPI_IOC_MAGIC: u8 = b'k'; // Defined in linux/spi/spidev.h
 // const SPI_IOC_TYPE_MESSAGE: u8 = 0;
@@ -41,29 +49,49 @@ ioctl_write_ptr!(
     DrmRockchipEbcOffScreen
 );
 
-fn main() -> eyre::Result<()> {
-    let opt = Opt::parse();
-    color_eyre::install()?;
+fn to_grayscale(path: &Path, image: DynamicImage) -> GrayImage {
+    if let DynamicImage::ImageLuma8(image) = image {
+        return image;
+    }
 
-    let mut image = ImageReader::open(&opt.path)?.decode()?;
+    eprintln!("Image {path:?} is not 8-bit grayscale. Converting...");
+    image.to_luma8()
+}
+
+fn load_image(path: &Path) -> eyre::Result<GrayImage> {
+    eprintln!("Opening image at {path:?}");
+
+    let image = ImageReader::open(path)?.decode()?;
+    Ok(to_grayscale(path, image))
+}
+
+fn load_wallpaper(path: &Path) -> eyre::Result<GrayImage> {
+    eprintln!("Opening image at {path:?}");
+
+    let mut image = ImageReader::open(path)?.decode()?;
 
     let image_dimensions = (image.width() as usize, image.height() as usize);
 
     match image_dimensions {
         (SCREEN_W, SCREEN_H) => {}
-        (SCREEN_H, SCREEN_W) => image = image.rotate90(),
+        (SCREEN_H, SCREEN_W) => image = image.rotate90().flipv(),
         _ => bail!("Image must be {SCREEN_W}x{SCREEN_H}"),
     }
 
-    let image = image.to_luma8();
-    let mut pixels = image.into_vec();
-    for pixel in &mut pixels {
-        *pixel >>= 4; // convert 8-bit colorspace to 4-bits.
-    }
-    pixels.push(0u8); // add a null-terminator
+    Ok(to_grayscale(path, image))
+}
 
-    // sanity check buffer length
-    assert!(pixels.len() == FILE_SIZE);
+fn main() -> eyre::Result<()> {
+    let opt = Opt::parse();
+    color_eyre::install()?;
+
+    let mut image = load_wallpaper(&opt.path)?;
+
+    for stamp in &opt.stamp {
+        apply_stamp(stamp, &mut image)?;
+    }
+
+    let pixels = image_to_offimage_buffer(image);
 
     let set_off_screen_data = DrmRockchipEbcOffScreen {
         info1: 0, // TODO: what is this?
@@ -80,3 +108,43 @@ fn main() -> eyre::Result<()> {
 
     Ok(())
 }
+
+fn image_to_offimage_buffer(image: image::ImageBuffer<image::Luma<u8>, Vec<u8>>) -> Vec<u8> {
+    let mut pixels = image.into_vec();
+
+    // Convert 8-bit colorspace to 4-bits by truncating the 4 least significant bits.
+    // Yes, this means that the most-significant 4 bits of each byte are unused.
+    // Not idea why they designed it this way. Maybe it's faster?
+    for pixel in &mut pixels {
+        *pixel >>= 4;
+    }
+
+    // Add a null-terminator. No idea why this is necessary, but it is.
+    pixels.push(0u8);
+
+    // Sanity check buffer length
+    assert!(pixels.len() == IOCTL_BUFFER_SIZE);
+
+    pixels
+}
+
+fn parse_stamp(stamp: &str) -> Option<(i64, i64, &Path)> {
+    let (x, stamp) = stamp.split_once(':')?;
+    let (y, path) = stamp.split_once(':')?;
+
+    let x = x.parse().ok()?;
+    let y = y.parse().ok()?;
+    let path = Path::new(path);
+    Some((x, y, path))
+}
+
+fn apply_stamp(stamp: &str, onto: &mut GrayImage) -> eyre::Result<()> {
+    let (x, y, path) =
+        parse_stamp(stamp).wrap_err_with(|| eyre!("Invalid stamp format: {stamp:?}"))?;
+
+    let image = load_image(path)?;
+
+    overlay(onto, &image, x, y);
+
+    Ok(())
+}