cyme 3.0.1

//! Leverages the usb HotplugEvent to create a stream of system USB devices
//!
//! See the watch cli for a usage example.
use super::nusb::NusbProfiler;
use super::{Device, DeviceEvent, SystemProfile};
use crate::error::Error;
use ::nusb::hotplug::HotplugEvent;
use ::nusb::watch_devices;
use chrono::Local;
use futures_lite::stream::Stream;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};

/// Builder for [`SystemProfileStream`]
#[derive(Default)]
pub struct SystemProfileStreamBuilder {
    spusb: Option<SystemProfile>,
    options: super::ProfilerOptions,
}

impl SystemProfileStreamBuilder {
    /// Create a new [`SystemProfileStreamBuilder`]
    pub fn new() -> Self {
        Self {
            spusb: None,
            // Default full and as tree for max information
            options: super::ProfilerOptions {
                filter: None,
                depth: super::ProfileDepth::Full,
                tree: true,
            },
        }
    }

    /// Set the verbosity of the stream
    ///
    /// When set to true, the stream will include full device descriptors for verbose printing
    pub fn is_verbose(mut self, verbose: bool) -> Self {
        self.options.set_verbose(verbose);
        self
    }

    /// Set the initial [`SystemProfile`] for the stream
    pub fn with_spusb(mut self, spusb: SystemProfile) -> Self {
        self.spusb = Some(spusb);
        self
    }

    /// Set the [`ProfilerOptions`] for the stream
    pub fn with_options(mut self, options: super::ProfilerOptions) -> Self {
        self.options = options;
        self
    }

    /// Build the [`SystemProfileStream`]
    pub fn build(self) -> Result<SystemProfileStream, Error> {
        let spusb = if let Some(spusb) = self.spusb {
            Arc::new(Mutex::new(spusb))
        } else {
            Arc::new(Mutex::new(super::get_spusb_with_options(&self.options)?))
        };

        SystemProfileStream::new_with_options(spusb, self.options)
    }
}

/// A stream that yields an updated [`SystemProfile`] when a USB device is connected or disconnected
pub struct SystemProfileStream {
    spusb: Arc<Mutex<SystemProfile>>,
    watch_stream: Pin<Box<dyn Stream<Item = HotplugEvent> + Send>>,
    options: super::ProfilerOptions,
}

impl SystemProfileStream {
    /// Create a new [`SystemProfileStream`] with a initial [`SystemProfile`]
    pub fn new(spusb: Arc<Mutex<SystemProfile>>) -> Result<Self, Error> {
        let watch_stream = Box::pin(watch_devices()?);
        Ok(Self {
            spusb,
            watch_stream,
            options: super::ProfilerOptions::default(),
        })
    }

    /// Create a new [`SystemProfileStream`] with an initial [`SystemProfile`] and [`ProfilerOptions`] for profiling devices
    pub fn new_with_options(
        spusb: Arc<Mutex<SystemProfile>>,
        options: super::ProfilerOptions,
    ) -> Result<Self, Error> {
        let watch_stream = Box::pin(watch_devices()?);
        Ok(Self {
            spusb,
            watch_stream,
            options,
        })
    }

    /// Get the [`SystemProfile`] from the stream
    pub fn get_profile(&self) -> Arc<Mutex<SystemProfile>> {
        Arc::clone(&self.spusb)
    }

    /// Re-profile the system USB devices
    ///
    /// Last events will be lost
    pub fn reprofile(&self) -> Arc<Mutex<SystemProfile>> {
        Arc::new(Mutex::new(
            super::get_spusb_with_options(&self.options).unwrap(),
        ))
    }
}

impl Stream for SystemProfileStream {
    type Item = Arc<Mutex<SystemProfile>>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.get_mut();
        let mut profiler = NusbProfiler::new();

        match Pin::new(&mut this.watch_stream).poll_next(cx) {
            Poll::Ready(Some(event)) => {
                let mut spusb = this.spusb.lock().unwrap();

                match event {
                    HotplugEvent::Connected(device) => {
                        let mut cyme_device: Device =
                            profiler.build_spdevice(&device, &this.options).unwrap();
                        cyme_device.last_event = Some(DeviceEvent::Connected(Local::now()));
                        // Windows bus number is a string ID so we need to find the bus based on this and assign the bus number created during cyme profiling
                        #[cfg(target_os = "windows")]
                        {
                            if let Some(existing_bus) =
                                spusb.buses.iter().find(|b| b.id == device.bus_id())
                            {
                                log::debug!(
                                    "Win found bus for connected device ({cyme_device}): {0}",
                                    existing_bus.id
                                );
                                if let Some(existing_number) = existing_bus.get_bus_number() {
                                    log::debug!(
                                        "Assigning bus number {existing_number} to device {cyme_device}",
                                    );
                                    cyme_device.location_id.bus = existing_number;
                                }
                            } else {
                                log::error!(
                                    "Win no bus found for connected device, seeking bus_id: {}",
                                    device.bus_id()
                                );
                            }
                        }

                        spusb.insert(cyme_device);
                    }
                    HotplugEvent::Disconnected(id) => {
                        if let Some(device) = spusb.get_id_mut(&id) {
                            device.last_event = Some(DeviceEvent::Disconnected(Local::now()));
                        }
                    }
                }
                Poll::Ready(Some(Arc::clone(&this.spusb)))
            }
            Poll::Ready(None) => Poll::Ready(None),
            Poll::Pending => Poll::Pending,
        }
    }
}