use std::{borrow::Borrow, collections::HashMap, sync::Arc, time::Duration}; use paris::warn; use tauri::async_runtime::RwLock; use tokio::{ net::UdpSocket, sync::{broadcast, watch}, time::sleep, }; use crate::{ ambient_light::{config, ConfigManager}, led_color::LedColor, rpc::UdpRpc, screenshot::{self, LedSamplePoints}, screenshot_manager::{self, ScreenshotManager}, }; use itertools::Itertools; use super::{LedStripConfigGroup, SamplePointMapper}; pub struct LedColorsPublisher { sorted_colors_rx: Arc>>>, sorted_colors_tx: Arc>>>, colors_rx: Arc>>>, colors_tx: Arc>>>, inner_tasks_version: Arc>, } impl LedColorsPublisher { pub async fn global() -> &'static Self { static LED_COLORS_PUBLISHER_GLOBAL: tokio::sync::OnceCell = tokio::sync::OnceCell::const_new(); let (sorted_tx, sorted_rx) = watch::channel(Vec::new()); let (tx, rx) = watch::channel(Vec::new()); LED_COLORS_PUBLISHER_GLOBAL .get_or_init(|| async { LedColorsPublisher { sorted_colors_rx: Arc::new(RwLock::new(sorted_rx)), sorted_colors_tx: Arc::new(RwLock::new(sorted_tx)), colors_rx: Arc::new(RwLock::new(rx)), colors_tx: Arc::new(RwLock::new(tx)), inner_tasks_version: Arc::new(RwLock::new(0)), } }) .await } async fn start_one_display_colors_fetcher( &self, display_id: u32, sample_points: Vec, bound_scale_factor: f32, mappers: Vec, display_colors_tx: broadcast::Sender<(u32, Vec)>, ) { let internal_tasks_version = self.inner_tasks_version.clone(); let screenshot_manager = ScreenshotManager::global().await; let screenshot_rx = screenshot_manager.subscribe_by_display_id(display_id).await; if let Err(err) = screenshot_rx { log::error!("{}", err); return; } let mut screenshot_rx = screenshot_rx.unwrap(); tokio::spawn(async move { let init_version = internal_tasks_version.read().await.clone(); while screenshot_rx.changed().await.is_ok() { let screenshot = screenshot_rx.borrow().clone(); let colors = screenshot.get_colors_by_sample_points(&sample_points).await; let colors_copy = colors.clone(); let mappers = mappers.clone(); match Self::send_colors_by_display(colors, mappers).await { Ok(_) => { // log::info!("sent colors: #{: >15}", display_id); } Err(err) => { warn!("Failed to send colors: #{: >15}\t{}", display_id, err); } } match display_colors_tx.send(( display_id, colors_copy .into_iter() .map(|color| color.get_rgb()) .flatten() .collect::>(), )) { Ok(_) => { // log::info!("sent colors: {:?}", color_len); } Err(err) => { warn!("Failed to send display_colors: {}", err); } }; // Check if the inner task version changed let version = internal_tasks_version.read().await.clone(); if version != init_version { log::info!( "inner task version changed, stop. {} != {}", internal_tasks_version.read().await.clone(), init_version ); break; } } }); } fn start_all_colors_worker( &self, display_ids: Vec, mappers: Vec, mut display_colors_rx: broadcast::Receiver<(u32, Vec)>, ) { let sorted_colors_tx = self.sorted_colors_tx.clone(); let colors_tx = self.colors_tx.clone(); tokio::spawn(async move { for _ in 0..10 { let sorted_colors_tx = sorted_colors_tx.write().await; let colors_tx = colors_tx.write().await; let mut all_colors: Vec>> = vec![None; display_ids.len()]; let mut start: tokio::time::Instant = tokio::time::Instant::now(); loop { let color_info = display_colors_rx.recv().await; if let Err(err) = color_info { match err { broadcast::error::RecvError::Closed => { return; } broadcast::error::RecvError::Lagged(_) => { warn!("display_colors_rx lagged"); continue; } } } let (display_id, colors) = color_info.unwrap(); let index = display_ids.iter().position(|id| *id == display_id); if index.is_none() { warn!("display id not found"); continue; } all_colors[index.unwrap()] = Some(colors); if all_colors.iter().all(|color| color.is_some()) { let flatten_colors = all_colors .clone() .into_iter() .flat_map(|c| c.unwrap()) .collect::>(); match colors_tx.send(flatten_colors.clone()) { Ok(_) => {} Err(err) => { warn!("Failed to send colors: {}", err); } }; let sorted_colors = ScreenshotManager::get_sorted_colors(&flatten_colors, &mappers); match sorted_colors_tx.send(sorted_colors) { Ok(_) => {} Err(err) => { warn!("Failed to send sorted colors: {}", err); } }; start = tokio::time::Instant::now(); } } } }); } pub async fn start(&self) { let config_manager = ConfigManager::global().await; let mut config_receiver = config_manager.clone_config_update_receiver(); let configs = config_receiver.borrow().clone(); self.handle_config_change(configs).await; while config_receiver.changed().await.is_ok() { let configs = config_receiver.borrow().clone(); self.handle_config_change(configs).await; } } async fn handle_config_change(&self, configs: LedStripConfigGroup) { let inner_tasks_version = self.inner_tasks_version.clone(); let configs = Self::get_colors_configs(&configs).await; if let Err(err) = configs { warn!("Failed to get configs: {}", err); sleep(Duration::from_millis(100)).await; return; } let configs = configs.unwrap(); let mut inner_tasks_version = inner_tasks_version.write().await; *inner_tasks_version = inner_tasks_version.overflowing_add(1).0; drop(inner_tasks_version); let (display_colors_tx, display_colors_rx) = broadcast::channel::<(u32, Vec)>(8); for sample_point_group in configs.sample_point_groups.clone() { let display_id = sample_point_group.display_id; let sample_points = sample_point_group.points; let bound_scale_factor = sample_point_group.bound_scale_factor; self.start_one_display_colors_fetcher( display_id, sample_points, bound_scale_factor, sample_point_group.mappers, display_colors_tx.clone(), ) .await; } let display_ids = configs.sample_point_groups; self.start_all_colors_worker( display_ids.iter().map(|c| c.display_id).collect(), configs.mappers, display_colors_rx, ); } pub async fn send_colors(offset: u16, mut payload: Vec) -> anyhow::Result<()> { // let mqtt = MqttRpc::global().await; // mqtt.publish_led_sub_pixels(payload).await; let socket = UdpSocket::bind("0.0.0.0:8000").await?; let mut buffer = vec![2]; buffer.push((offset >> 8) as u8); buffer.push((offset & 0xff) as u8); buffer.append(&mut payload); socket.send_to(&buffer, "192.168.31.206:23042").await?; Ok(()) } pub async fn send_colors_by_display( colors: Vec, mappers: Vec, ) -> anyhow::Result<()> { // let color_len = colors.len(); let display_led_offset = mappers .clone() .iter() .flat_map(|mapper| [mapper.start, mapper.end]) .min() .unwrap(); let udp_rpc = UdpRpc::global().await; if let Err(err) = udp_rpc { warn!("udp_rpc can not be initialized: {}", err); } let udp_rpc = udp_rpc.as_ref().unwrap(); // let socket = UdpSocket::bind("0.0.0.0:0").await?; for group in mappers.clone() { if (group.start.abs_diff(group.end)) > colors.len() { return Err(anyhow::anyhow!( "get_sorted_colors: color_index out of range. color_index: {}, strip len: {}, colors.len(): {}", group.pos, group.start.abs_diff(group.end), colors.len() )); } let group_size = group.start.abs_diff(group.end); let mut buffer = Vec::::with_capacity(group_size * 3); if group.end > group.start { // Prevent integer underflow by using saturating subtraction let start_index = if group.pos >= display_led_offset { group.pos - display_led_offset } else { 0 }; let end_index = if group.pos + group_size >= display_led_offset { group_size + group.pos - display_led_offset } else { 0 }; for i in start_index..end_index { if i < colors.len() { let bytes = colors[i].as_bytes(); buffer.append(&mut bytes.to_vec()); } else { log::warn!("Index {} out of bounds for colors array of length {}", i, colors.len()); // Add black color as fallback buffer.append(&mut vec![0, 0, 0]); } } } else { // Prevent integer underflow by using saturating subtraction let start_index = if group.pos >= display_led_offset { group.pos - display_led_offset } else { 0 }; let end_index = if group.pos + group_size >= display_led_offset { group_size + group.pos - display_led_offset } else { 0 }; for i in (start_index..end_index).rev() { if i < colors.len() { let bytes = colors[i].as_bytes(); buffer.append(&mut bytes.to_vec()); } else { log::warn!("Index {} out of bounds for colors array of length {}", i, colors.len()); // Add black color as fallback buffer.append(&mut vec![0, 0, 0]); } } } let offset = group.start.min(group.end); let mut tx_buffer = vec![2]; tx_buffer.push((offset >> 8) as u8); tx_buffer.push((offset & 0xff) as u8); tx_buffer.append(&mut buffer); udp_rpc.send_to_all(&tx_buffer).await?; } Ok(()) } pub async fn clone_sorted_colors_receiver(&self) -> watch::Receiver> { self.sorted_colors_rx.read().await.clone() } pub async fn get_colors_configs( configs: &LedStripConfigGroup, ) -> anyhow::Result { let screenshot_manager = ScreenshotManager::global().await; let display_ids = configs .strips .iter() .map(|c| c.display_id) .unique() .collect::>(); let mappers = configs.mappers.clone(); let mut colors_configs = Vec::new(); let mut merged_screenshot_receiver = screenshot_manager.clone_merged_screenshot_rx().await; merged_screenshot_receiver.resubscribe(); let mut screenshots = HashMap::new(); loop { let screenshot = merged_screenshot_receiver.recv().await; if let Err(err) = screenshot { match err { tokio::sync::broadcast::error::RecvError::Closed => { warn!("closed"); continue; } tokio::sync::broadcast::error::RecvError::Lagged(_) => { warn!("lagged"); continue; } } } let screenshot = screenshot.unwrap(); // log::info!("got screenshot: {:?}", screenshot.display_id); screenshots.insert(screenshot.display_id, screenshot); if screenshots.len() == display_ids.len() { let mut led_start = 0; for display_id in display_ids { let led_strip_configs = configs .strips .iter() .enumerate() .filter(|(_, c)| c.display_id == display_id); let screenshot = screenshots.get(&display_id).unwrap(); let points: Vec<_> = led_strip_configs .clone() .map(|(_, config)| screenshot.get_sample_points(&config)) .flatten() .collect(); if points.len() == 0 { warn!("no led strip config for display_id: {}", display_id); continue; } let bound_scale_factor = screenshot.bound_scale_factor; let led_end = led_start + points.iter().map(|p| p.len()).sum::(); let mappers = led_strip_configs.map(|(i, _)| mappers[i].clone()).collect(); let colors_config = DisplaySamplePointGroup { display_id, points, bound_scale_factor, mappers, }; colors_configs.push(colors_config); led_start = led_end; } return Ok(AllColorConfig { sample_point_groups: colors_configs, mappers, }); } } } pub async fn clone_colors_receiver(&self) -> watch::Receiver> { self.colors_rx.read().await.clone() } } #[derive(Debug)] pub struct AllColorConfig { pub sample_point_groups: Vec, pub mappers: Vec, // pub screenshot_receivers: Vec>, } #[derive(Debug, Clone)] pub struct DisplaySamplePointGroup { pub display_id: u32, pub points: Vec, pub bound_scale_factor: f32, pub mappers: Vec, }