desktop/src-tauri/src/picker/manager.rs

use futures::{stream::FuturesUnordered, StreamExt};
use once_cell::sync::OnceCell;
use paris::info;
use scrap::Display;
use std::sync::Arc;
use tokio::{sync::Mutex, task};

use crate::picker::{config::LedStripConfig, screen::Screen};

use super::{
    config::DisplayConfig, display_picker::DisplayPicker, led_color::LedColor,
    screenshot::Screenshot,
};

pub struct Picker {
    pub screens: Arc<Mutex<Vec<Screen>>>,
    pub screenshots: Arc<Mutex<Vec<Screenshot>>>,
    pub display_configs: Arc<Mutex<Vec<DisplayConfig>>>,
}

impl Picker {
    pub fn global() -> &'static Picker {
        static SCREEN_COLOR_PICKER: OnceCell<Picker> = OnceCell::new();

        SCREEN_COLOR_PICKER.get_or_init(|| Picker {
            screens: Arc::new(Mutex::new(vec![])),
            screenshots: Arc::new(Mutex::new(vec![])),
            display_configs: Arc::new(Mutex::new(vec![
                DisplayConfig {
                    index_of_display: 1,
                    display_width: 1920,
                    display_height: 1200,
                    top_led_strip: LedStripConfig {
                        index: 1,
                        global_start_position: 59,
                        global_end_position: 32,
                    },
                    bottom_led_strip: LedStripConfig {
                        index: 0,
                        global_start_position: 0,
                        global_end_position: 0,
                    },
                    left_led_strip: LedStripConfig {
                        index: 0,
                        global_start_position: 0,
                        global_end_position: 0,
                    },
                    right_led_strip: LedStripConfig {
                        index: 0,
                        global_start_position: 0,
                        global_end_position: 0,
                    },
                },
                DisplayConfig {
                    index_of_display: 0,
                    display_width: 3008,
                    display_height: 1692,
                    top_led_strip: LedStripConfig {
                        index: 0,
                        global_start_position: 31,
                        global_end_position: 0,
                    },
                    bottom_led_strip: LedStripConfig {
                        index: 0,
                        global_start_position: 0,
                        global_end_position: 0,
                    },
                    left_led_strip: LedStripConfig {
                        index: 0,
                        global_start_position: 0,
                        global_end_position: 0,
                    },
                    right_led_strip: LedStripConfig {
                        index: 0,
                        global_start_position: 0,
                        global_end_position: 0,
                    },
                },
            ])),
        })
    }

    pub async fn list_displays(&self) -> anyhow::Result<Vec<String>> {
        let mut configs = self.display_configs.lock().await;
        let screenshots = self.screenshots.lock().await;

        let displays = Display::all()
            .map_err(|error| anyhow::anyhow!("Can not get all of displays. {}", error))?;

        configs.clear();
        let mut futs = FuturesUnordered::new();

        for (index, display) in displays.iter().enumerate() {
            let height = display.height();
            let width = display.width();
            let config = DisplayConfig::default(index, width, height);
            configs.push(config);
        }

        for (index, display) in displays.iter().enumerate() {
            let height = display.height();
            let width = display.width();
            let config = configs[index];
            futs.push(async move {
                let join = task::spawn(Self::preview_display_by_config(config));
                join.await?
            });
        }
        let mut bitmap_string_list = vec![];
        while let Some(bitmap_string) = futs.next().await {
            match bitmap_string {
                Ok(bitmap_string) => {
                    bitmap_string_list.push(bitmap_string);
                }
                Err(error) => {
                    anyhow::bail!("can not convert to base64 image. {}", error);
                }
            }
        }
        Ok(bitmap_string_list)
    }

    pub async fn preview_display_by_config(config: DisplayConfig) -> anyhow::Result<String> {
        let start = time::Instant::now();
        let mut picker = DisplayPicker::from_config(config)?;
        let screenshot = picker.take_screenshot()?;
        info!("Take Screenshot Spend: {}", start.elapsed());

        anyhow::Ok(screenshot.to_webp_base64().await)
    }

    pub async fn refresh_displays(&self) -> anyhow::Result<()> {
        // let displays = Display::all()
        //     .map_err(|error| anyhow::anyhow!("Can not get all of displays. {}", error))?;
        // let mut screens = self.screens.lock().await;
        // let mut screenshots = self.screenshots.lock().await;
        // screens.clear();
        // info!("number of displays: {}", displays.len());
        // for display in displays {
        //     let height = display.height();
        //     let width = display.width();
        //     match Capturer::new(display) {
        //         Ok(capturer) => screens.push(Screen::new(capturer, width, height)),
        //         Err(error) => screens.push(Screen::new_failed(
        //             anyhow::anyhow!("{}", error),
        //             width,
        //             height,
        //         )),
        //     };
        //     screenshots.push(Screenshot::new(width, height));
        // }

        // screens.reverse();
        // screenshots.reverse();

        Ok(())
    }

    pub async fn take_screenshots_for_all(&self) -> anyhow::Result<Vec<Screenshot>> {
        let mut screens = self.screens.lock().await;
        let screenshots = self.screenshots.lock().await;
        for (index, screen) in screens.iter_mut().enumerate() {
            let bitmap = screen.take().map_err(|error| {
                anyhow::anyhow!("take screenshot for display failed. {}", error)
            })?;
        }
        Ok(screenshots.to_vec())
    }

    pub async fn get_led_strip_colors(&self) -> anyhow::Result<Vec<LedColor>> {
        let screenshots = self.screenshots.lock().await;
        let mut colors = Vec::new();
        for screenshot in screenshots.iter() {
            let result = screenshot
                .get_top_colors()
                .await
                .map_err(|error| anyhow::anyhow!("get top colors failed. {}", error))?;
            colors.extend_from_slice(&result);
        }
        Ok(colors)
    }
}