automation_rs/src/devices.rs

mod audio_setup;
mod contact_sensor;
mod ikea_outlet;
mod kasa_outlet;
mod wake_on_lan;

pub use self::audio_setup::AudioSetup;
pub use self::contact_sensor::ContactSensor;
pub use self::ikea_outlet::IkeaOutlet;
pub use self::kasa_outlet::KasaOutlet;
pub use self::wake_on_lan::WakeOnLAN;

use std::collections::HashMap;

use google_home::{traits::OnOff, FullfillmentError, GoogleHome, GoogleHomeDevice};
use pollster::FutureExt;
use rumqttc::{matches, AsyncClient, QoS};
use thiserror::Error;
use tokio::sync::{mpsc, oneshot};
use tracing::{debug, error, trace};

use crate::{
    event::{Event, EventChannel},
    light_sensor::OnDarkness,
    mqtt::OnMqtt,
    presence::OnPresence,
};

#[impl_cast::device(As: OnMqtt + OnPresence + OnDarkness + GoogleHomeDevice + OnOff)]
pub trait Device: std::fmt::Debug + Sync + Send {
    fn get_id(&self) -> &str;
}

// TODO: Add an inner type that we can wrap with Arc<RwLock<>> to make this type a little bit nicer
// to work with
#[derive(Debug)]
struct Devices {
    devices: HashMap<String, Box<dyn Device>>,
    client: AsyncClient,
}

#[derive(Debug)]
pub enum Command {
    Fullfillment {
        google_home: GoogleHome,
        payload: google_home::Request,
        tx: oneshot::Sender<Result<google_home::Response, FullfillmentError>>,
    },
    AddDevice {
        device: Box<dyn Device>,
        tx: oneshot::Sender<()>,
    },
}

#[derive(Clone)]
pub struct DevicesHandle {
    tx: mpsc::Sender<Command>,
}

#[derive(Debug, Error)]
pub enum DevicesError {
    #[error(transparent)]
    FullfillmentError(#[from] FullfillmentError),
    #[error(transparent)]
    SendError(#[from] tokio::sync::mpsc::error::SendError<Command>),
    #[error(transparent)]
    RecvError(#[from] tokio::sync::oneshot::error::RecvError),
}

impl DevicesHandle {
    // TODO: Improve error type
    pub async fn fullfillment(
        &self,
        google_home: GoogleHome,
        payload: google_home::Request,
    ) -> Result<google_home::Response, DevicesError> {
        let (tx, rx) = oneshot::channel();
        self.tx
            .send(Command::Fullfillment {
                google_home,
                payload,
                tx,
            })
            .await?;
        Ok(rx.await??)
    }

    pub async fn add_device(&self, device: Box<dyn Device>) -> Result<(), DevicesError> {
        let (tx, rx) = oneshot::channel();
        self.tx.send(Command::AddDevice { device, tx }).await?;
        Ok(rx.await?)
    }
}

pub fn start(event_channel: &EventChannel, client: AsyncClient) -> DevicesHandle {
    let mut devices = Devices {
        devices: HashMap::new(),
        client,
    };

    let (tx, mut rx) = mpsc::channel(100);
    let mut event_rx = event_channel.get_rx();

    tokio::spawn(async move {
        // TODO: Handle error better
        loop {
            tokio::select! {
                event = event_rx.recv() => {
                    if event.is_err() {
                        todo!("Handle errors with the event channel properly")
                    }
                    devices.handle_event(event.unwrap()).await;
                }
                // TODO: Handle receiving None better, otherwise it might constantly run doing
                // nothing
                cmd = rx.recv() => {
                    if cmd.is_none() {
                        todo!("Handle errors with the cmd channel properly")
                    }
                    devices.handle_cmd(cmd.unwrap()).await;
                }
            }
        }
    });

    DevicesHandle { tx }
}

impl Devices {
    async fn handle_cmd(&mut self, cmd: Command) {
        match cmd {
            Command::Fullfillment {
                google_home,
                payload,
                tx,
            } => {
                let result =
                    google_home.handle_request(payload, &mut self.get::<dyn GoogleHomeDevice>());
                tx.send(result).ok();
            }
            Command::AddDevice { device, tx } => {
                self.add_device(device).await;

                tx.send(()).ok();
            }
        }
    }

    async fn add_device(&mut self, device: Box<dyn Device>) {
        let id = device.get_id();
        debug!(id, "Adding device");

        // If the device listens to mqtt, subscribe to the topics
        if let Some(device) = As::<dyn OnMqtt>::cast(device.as_ref()) {
            for topic in device.topics() {
                trace!(id, topic, "Subscribing to topic");
                if let Err(err) = self.client.subscribe(topic, QoS::AtLeastOnce).await {
                    // NOTE: Pretty sure that this can only happen if the mqtt client if no longer
                    // running
                    error!(id, topic, "Failed to subscribe to topic: {err}");
                }
            }
        }

        self.devices.insert(device.get_id().to_owned(), device);
    }

    async fn handle_event(&mut self, event: Event) {
        match event {
            Event::MqttMessage(message) => {
                self.get::<dyn OnMqtt>()
                    .iter_mut()
                    .for_each(|(id, listener)| {
                        let subscribed = listener
                            .topics()
                            .iter()
                            .any(|topic| matches(&message.topic, topic));

                        if subscribed {
                            trace!(id, "Handling");
                            listener.on_mqtt(&message).block_on();
                        }
                    })
            }
            Event::Darkness(dark) => {
                self.get::<dyn OnDarkness>()
                    .iter_mut()
                    .for_each(|(id, device)| {
                        trace!(id, "Handling");
                        device.on_darkness(dark).block_on();
                    })
            }
            Event::Presence(presence) => {
                self.get::<dyn OnPresence>()
                    .iter_mut()
                    .for_each(|(id, device)| {
                        trace!(id, "Handling");
                        device.on_presence(presence).block_on();
                    })
            }
            Event::Ntfy(_) => {}
        }
    }

    fn get<T>(&mut self) -> HashMap<&str, &mut T>
    where
        T: ?Sized + 'static,
        (dyn Device): As<T>,
    {
        self.devices
            .iter_mut()
            .filter_map(|(id, device)| As::<T>::cast_mut(device.as_mut()).map(|t| (id.as_str(), t)))
            .collect()
    }
}