use async_trait::async_trait; use image::{AnimationDecoder, DynamicImage, codecs::gif::GifDecoder, imageops::FilterType}; use std::io::Cursor; use std::time::Duration; use crate::config; use crate::display::DisplayState; use rayhunter::analysis::analyzer::EventType; use log::{error, info}; use tokio::sync::mpsc::Receiver; use tokio_util::{sync::CancellationToken, task::TaskTracker}; use include_dir::{Dir, include_dir}; const REFRESH_RATE: u64 = 1000; //how often in milliseconds to refresh the display #[derive(Copy, Clone)] pub struct Dimensions { pub height: u32, pub width: u32, } #[derive(Copy, Clone)] pub enum LinePattern { Solid, Dashed, // _ _ _ _ Dotted, // . . . . } #[allow(dead_code)] #[derive(Copy, Clone)] pub enum Color { Red, Green, Blue, White, Black, Cyan, Yellow, Pink, Orange, } impl Color { fn rgb(self) -> (u8, u8, u8) { match self { Color::Red => (0xff, 0, 0), Color::Green => (0, 0xff, 0), Color::Blue => (0, 0, 0xff), Color::White => (0xff, 0xff, 0xff), Color::Black => (0, 0, 0), Color::Cyan => (0, 0xff, 0xff), Color::Yellow => (0xff, 0xff, 0), Color::Pink => (0xfe, 0x24, 0xff), Color::Orange => (0xff, 0xa5, 0), } } } fn display_style_from_state(state: DisplayState, colorblind_mode: bool) -> (Color, LinePattern) { match state { DisplayState::Paused => (Color::White, LinePattern::Solid), DisplayState::Recording => { if colorblind_mode { (Color::Blue, LinePattern::Solid) } else { (Color::Green, LinePattern::Solid) } } DisplayState::WarningDetected { event_type } => match event_type { EventType::Informational => { if colorblind_mode { (Color::Blue, LinePattern::Solid) } else { (Color::Green, LinePattern::Solid) } } EventType::Low => (Color::Yellow, LinePattern::Dotted), EventType::Medium => (Color::Orange, LinePattern::Dashed), EventType::High => (Color::Red, LinePattern::Solid), }, } } #[async_trait] pub trait GenericFramebuffer: Send + 'static { fn dimensions(&self) -> Dimensions; async fn write_buffer(&mut self, buffer: Vec<(u8, u8, u8)>); // rgb, row-wise, left-to-right, top-to-bottom async fn write_dynamic_image(&mut self, img: DynamicImage) { let dimensions = self.dimensions(); let mut width = img.width(); let mut height = img.height(); let resized_img: DynamicImage; if height > dimensions.height || width > dimensions.width { resized_img = img.resize(dimensions.width, dimensions.height, FilterType::CatmullRom); width = dimensions.width.min(resized_img.width()); height = dimensions.height.min(resized_img.height()); } else { resized_img = img; } let img_rgba8 = resized_img.as_rgba8().unwrap(); let mut buf = Vec::new(); for y in 0..height { for x in 0..width { let px = img_rgba8.get_pixel(x, y); buf.push((px[0], px[1], px[2])); } } self.write_buffer(buf).await } async fn draw_gif(&mut self, img_buffer: &[u8]) { let cursor = Cursor::new(img_buffer); if let Ok(decoder) = GifDecoder::new(cursor) { let frames: Vec<_> = decoder .into_frames() .filter_map(|f| f.ok()) .map(|frame| { let (numerator, _) = frame.delay().numer_denom_ms(); let img = DynamicImage::from(frame.into_buffer()); (img, numerator as u64) }) .collect(); for (img, delay_ms) in frames { self.write_dynamic_image(img).await; tokio::time::sleep(Duration::from_millis(delay_ms)).await; } } } async fn draw_img(&mut self, img_buffer: &[u8]) { let img = image::load_from_memory(img_buffer).unwrap(); self.write_dynamic_image(img).await } async fn draw_line(&mut self, color: Color, height: u32) { self.draw_patterned_line(color, height, LinePattern::Solid) .await } async fn draw_patterned_line(&mut self, color: Color, height: u32, pattern: LinePattern) { let width = self.dimensions().width; let mut buffer = Vec::new(); for _row in 0..height { for col in 0..width { let should_draw = match pattern { LinePattern::Solid => true, LinePattern::Dashed => (col / 4) % 2 == 0, // 4 pixels on, 4 pixels off LinePattern::Dotted => col % 4 == 0, // 1 pixel on, 3 pixels off }; if should_draw { buffer.push(color.rgb()); } else { buffer.push((0, 0, 0)); // Black background } } } self.write_buffer(buffer).await } } pub fn update_ui( task_tracker: &TaskTracker, config: &config::Config, mut fb: impl GenericFramebuffer, shutdown_token: CancellationToken, mut ui_update_rx: Receiver, ) { static IMAGE_DIR: Dir<'_> = include_dir!("$CARGO_MANIFEST_DIR/images/"); let display_level = config.ui_level; if display_level == 0 { info!("Invisible mode, not spawning UI."); } let colorblind_mode = config.colorblind_mode; let mut display_style = display_style_from_state(DisplayState::Recording, colorblind_mode); task_tracker.spawn(async move { // this feels wrong, is there a more rusty way to do this? let mut img: Option<&[u8]> = None; if display_level == 2 { img = Some( IMAGE_DIR .get_file("orca.gif") .expect("failed to read orca.gif") .contents(), ); } else if display_level == 3 { img = Some( IMAGE_DIR .get_file("eff.png") .expect("failed to read eff.png") .contents(), ); } loop { if shutdown_token.is_cancelled() { info!("received UI shutdown"); break; } match ui_update_rx.try_recv() { Ok(state) => { display_style = display_style_from_state(state, colorblind_mode); } Err(tokio::sync::mpsc::error::TryRecvError::Empty) => {} Err(e) => error!("error receiving framebuffer update message: {e}"), } match display_level { 2 => fb.draw_gif(img.unwrap()).await, 3 => fb.draw_img(img.unwrap()).await, 128 => { fb.draw_line(Color::Cyan, 128).await; fb.draw_line(Color::Pink, 102).await; fb.draw_line(Color::White, 76).await; fb.draw_line(Color::Pink, 50).await; fb.draw_line(Color::Cyan, 25).await; } // this branch id for ui_level 1, which is also the default if an // unknown value is used _ => {} }; let (color, pattern) = display_style; fb.draw_patterned_line(color, 2, pattern).await; tokio::time::sleep(Duration::from_millis(REFRESH_RATE)).await; } }); }