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

use once_cell::sync::OnceCell;
use paris::*;
use scrap::{Capturer, Display};
use std::sync::Arc;
use tokio::sync::Mutex;

use crate::picker::screen::Screen;

use super::{led_color::LedColor, screenshot::Screenshot};

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

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![])),
        })
    }

    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();
        screenshots[0].set_number_of_leds(22, 0);
        screenshots[1].set_number_of_leds(38, 0);

        Ok(())
    }

    pub async fn take_screenshots_for_all(&self) -> anyhow::Result<Vec<Screenshot>> {
        let mut screens = self.screens.lock().await;
        let mut 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)
            })?;
            screenshots[index].set_bitmap(bitmap).await
        }
        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)
    }
}