smart-home-in-rust-with-home-assistant/home-assistant/src/state_object.rs

use super::{
    event::{context::Context, specific::state_changed},
    home_assistant::HomeAssistant,
};
use crate::{entity_id::EntityId, home_assistant::GetStatesError, state_machine::GetStateError};
use chrono::{DateTime, Utc};
use emitter_and_signal::signal::Signal;
use once_cell::sync::OnceCell;
use pyo3::{
    types::{PyAnyMethods as _, PyCFunction, PyDict, PyModule, PyTuple},
    Bound, FromPyObject, IntoPyObject as _, Py, PyAny, PyErr, Python,
};
use snafu::{ResultExt, Snafu};
use std::{future::Future, sync::Arc};
use tokio::{select, sync::mpsc};

#[derive(Debug, FromPyObject)]
pub struct StateObject<State, Attributes, ContextEvent> {
    pub entity_id: EntityId,
    pub state: State,
    pub attributes: Attributes,
    pub last_changed: Option<DateTime<Utc>>,
    pub last_reported: Option<DateTime<Utc>>,
    pub last_updated: Option<DateTime<Utc>>,
    pub context: Context<ContextEvent>,
}

pub type ExtractStateObjectError<'a, 'py, State, Attributes, ContextEvent> =
    <StateObject<State, Attributes, ContextEvent> as FromPyObject<'a, 'py>>::Error;

#[derive(Debug, Snafu)]
pub enum CreateSignalError {
    /// couldn't get the state machine from the Home Assistant object
    GetStatesError { source: GetStatesError },
}

#[derive(Debug, Clone, Snafu)]
pub enum StateObjectSignalError<ExtractStateObjectError: 'static + snafu::Error> {
    /// couldn't get the state object directly from the state machine
    GetFromStateMachine {
        source: GetStateError<ExtractStateObjectError>,
    },

    /// couldn't get the state object from the new state event
    GetFromNewStateEvent { source: Arc<PyErr> },
}

impl<
        State: Send + Sync + 'static + for<'a, 'py> FromPyObject<'a, 'py>,
        Attributes: Send + Sync + 'static + for<'a, 'py> FromPyObject<'a, 'py>,
        ContextEvent: Send + Sync + 'static + for<'a, 'py> FromPyObject<'a, 'py>,
    > StateObject<State, Attributes, ContextEvent>
{
    pub fn signal<'a, 'py>(
        py: Python<'py>,
        home_assistant: &'py HomeAssistant,
        entity_id: EntityId,
    ) -> Result<
        (
            Signal<
                Option<
                    Arc<
                        Result<
                            Self,
                            StateObjectSignalError<Arc<<Self as FromPyObject<'a, 'py>>::Error>>,
                        >,
                    >,
                >,
            >,
            impl Future<Output = Result<(), emitter_and_signal::signal::JoinError>>,
        ),
        CreateSignalError,
    > {
        let state_machine = home_assistant.states(py).context(GetStatesSnafu)?;
        let current = state_machine
            .get(py, entity_id.clone())
            .map_err(|e| match e {
                GetStateError::GetStateObjectError { source } => {
                    GetStateError::GetStateObjectError { source }
                }
                GetStateError::ExtractStateObjectError { source } => {
                    GetStateError::ExtractStateObjectError {
                        source: Arc::new(source),
                    }
                }
            })
            .context(GetFromStateMachineSnafu)
            .transpose();

        let Ok(py_home_assistant) = home_assistant.into_pyobject(py);
        let py_home_assistant = py_home_assistant.unbind();

        let (signal, task) = Signal::new(
            current.map(Arc::new),
            |mut publisher_stream| async move {
                while let Some(publisher) = publisher_stream.wait().await {
                    let (new_state_sender, mut new_state_receiver) = mpsc::channel(8);

                    let untrack = Python::attach::<_, Result<_, PyErr>>(|py| {
                        static EVENT_MODULE: OnceCell<Py<PyModule>> = OnceCell::new();

                        let event_module = EVENT_MODULE
                            .get_or_try_init(|| {
                                Result::<_, PyErr>::Ok(
                                    py.import("homeassistant.helpers.event")?.unbind(),
                                )
                            })?
                            .bind(py);

                        let untrack = {
                            let callback =
                            move |args: &Bound<'_, PyTuple>,
                                  _kwargs: Option<&Bound<'_, PyDict>>| {
                                #[cfg(feature = "tracing")]
                                tracing::debug!("calling the closure");

                                let new_state_res = args.extract::<(
                                    state_changed::Event<
                                        State,
                                        Attributes,
                                        ContextEvent,
                                        State,
                                        Attributes,
                                        ContextEvent,
                                        Py<PyAny>,
                                    >,
                                )>().map(|event| event.0.data.new_state).map_err(Arc::new).context(GetFromNewStateEventSnafu);

                                new_state_sender.try_send(new_state_res).unwrap();
                            };
                            let callback = PyCFunction::new_closure(py, None, None, callback)?;
                            let args = (
                                py_home_assistant.clone_ref(py),
                                vec![entity_id.clone()],
                                callback,
                            );
                            event_module.call_method1("async_track_state_change_event", args)?
                        };

                        let untrack = untrack.unbind();

                        Ok(untrack)
                    });

                    if let Ok(untrack) = untrack {
                        #[cfg(feature = "tracing")]
                        tracing::debug!(
                            "untrack is ok, going to wait for the next relevant event..."
                        );
                        loop {
                            select! {
                                biased;
                                _ = publisher.all_unsubscribed() => {
                                    #[cfg(feature = "tracing")]
                                    tracing::debug!("calling untrack");
                                    let res = Python::attach(|py| untrack.call0(py));

                                    #[cfg(feature = "tracing")]
                                    tracing::debug!(?res);
                                    break;
                                }
                                new_state_res_option = new_state_receiver.recv() => {
                                    match new_state_res_option {
                                        Some(new_state_res) => {
                                            #[cfg(feature = "tracing")]
                                            tracing::debug!("publishing new state");
                                            publisher.publish(new_state_res.transpose().map(Arc::new));
                                        },
                                        None => {
                                            #[cfg(feature = "tracing")]
                                            tracing::debug!("channel dropped");
                                            break
                                        },
                                    }
                                }
                            }
                        }
                    } else {
                        #[cfg(feature = "tracing")]
                        tracing::debug!("untrack is err");
                    }
                }
            },
        );

        Ok((signal, task))
    }
}