mechutil 0.8.1

//
// Copyright (C) 2024 Automated Design Corp.. All Rights Reserved.
// Created Date: 2024-10-19 17:35:54
// -----
// Last Modified: 2025-03-22 09:10:15
// -----
//
//

//! This module provides an asynchronous implementation of an IBus (Inter-Bus) node, which is a component
//! that can communicate with other nodes in a distributed system. The node can handle requests, broadcasts,
//! subscriptions, and control messages, and it operates in different states (Init, Config, PreOp, Run).
//!
//! ## Example
//!
//! ```rust
//! use async_ibus_node::{AsyncIBusNode, AsyncIBusNodeFactory, AsyncIBusNodeHandle, AsyncIBusOpMode};
//! use serde_json::json;
//! use tokio::sync::mpsc;
//!
//! #[derive(Clone)]
//! struct MyNode {
//!     node_id: String,
//! }
//!
//! #[async_trait]
//! impl AsyncIBusNode<String, String> for MyNode {
//!     fn node_id(&self) -> String {
//!         self.node_id.clone()
//!     }
//!
//!     async fn request_change_opmode(&mut self, opmode: AsyncIBusOpMode, respond_to: oneshot::Sender<String>) {
//!         // Handle opmode change request
//!     }
//!
//!     async fn request_read_opmode(&self, respond_to: oneshot::Sender<String>) {
//!         // Handle opmode read request
//!     }
//!
//!     async fn request_received(&mut self, msg: String, respond_to: oneshot::Sender<String>) {
//!         // Handle request
//!     }
//!
//!     async fn broadcast_received(&mut self, msg: String) {
//!         // Handle broadcast
//!     }
//! }
//!
//! impl AsyncIBusNodeFactory<String, String> for MyNode {
//!     fn init_node(node_id: String, registry_handle: AsyncNodeRegistryHandle<String, String>, configuration: &serde_json::Value) -> Self {
//!         MyNode { node_id }
//!     }
//!
//! #[tokio::main]
//! async fn main() {
//!     let registry_handle = AsyncNodeRegistryHandle::new();
//!     let configuration = json!({});
//!     let node_handle = MyNode::create("my_node", registry_handle, &configuration).unwrap();
//!     // Use node_handle to interact with the node
//! }
//! ```

use anyhow::anyhow;
use async_trait::async_trait;
use log::debug;
use std::marker::PhantomData;
use tokio::sync::{mpsc, oneshot};

use super::async_node_registry::AsyncNodeRegistryHandle;

#[derive(Copy, Clone, Debug, PartialEq)]
#[repr(u8)]
pub enum AsyncIBusOpMode {
    /// Initialization State.
    /// The Node begins in this state and does nothing.
    /// The link to the registry is established, and the node receives is configuration
    /// information, but does nothing with it.
    Init = 0,
    /// Configuration State.
    /// The node is expected to load its configuration and make any
    /// internal preparations.
    Config = 2,
    /// Pre-Op State.
    /// The node makes its subscriptions or requests with other nodes. By waiting until this state,
    /// the other nodes in the network are ensured to be ready to handle subscriptions and requests.
    PreOp = 16,
    /// Run State.
    /// The Node is fully operational and ready to handle requests and perform its intended functions.
    Run = 32,
}

/// Checks if the transition from `current_opmode` to `requested_opmode` is valid.
///
/// # Parameters
/// - `current_opmode`: The current operation mode of the node.
/// - `requested_opmode`: The operation mode the node is attempting to transition to.
///
/// # Returns
/// - `bool`: `true` if the transition is valid, `false` otherwise.
pub fn is_valid_opmode_transition(
    current_opmode: AsyncIBusOpMode,
    requested_opmode: AsyncIBusOpMode,
) -> bool {
    match (current_opmode, requested_opmode) {
        // Valid transitions
        (AsyncIBusOpMode::Init, AsyncIBusOpMode::Config) => true,
        (AsyncIBusOpMode::Config, AsyncIBusOpMode::PreOp) => true,
        (AsyncIBusOpMode::PreOp, AsyncIBusOpMode::Run) => true,
        // Dropping down to Init is always valid
        (_, AsyncIBusOpMode::Init) => true,
        // Staying in the same state is also valid
        (current, requested) if current == requested => true,
        // All other transitions are invalid
        _ => false,
    }
}

/// Wraps a payload with the additional members we need for
/// standard messages on the Bus.
pub enum IBusMessage<S, R> {
    /// Make a request with message payload type S. The response will be
    /// received via a oneshot Sender with message of type R.
    Request {
        payload: S,
        respond_to: oneshot::Sender<R>,
    },
    /// A "broadcast" is an unsolicited message, similar to a Request but
    /// expecting no response. It contains only the payload of type R.
    Broadcast { payload: R },
    // /// Subscribe to messages from a module. Broadcasts will be received on
    // /// a broadcast channel that will be left open so long as the remote module
    // /// exists and it doesn't receive an unsubscribe request.
    // /// We use a broadcast channel instead of an mpsc so that a module instance
    // /// can centralize its receiver logic, if desired.
    // ///
    // /// The payload should be used for additional subscription details required
    // /// for a specific implementation.
    // Subscribe {
    //     topic : String,
    //     broadcast_tx: SubscriptionBroadcastTx<R>,
    //     respond_to: SubscriptionResponseTx,
    // },
    // /// Subscribe to messages from a module, using a filter to select the messages
    // /// received.
    // /// Broadcasts will be received on
    // /// a broadcast channel that will be left open so long as the remote module
    // /// exists and it doesn't receive an unsubscribe request.
    // /// We use a broadcast channel instead of an mpsc so that a module instance
    // /// can centralize its receiver logic, if desired.
    // ///
    // /// The payload should be used for additional subscription details required
    // /// for a specific implementation.
    // SubscribeWithFilter {
    //     topic : String,
    //     filter: Box<dyn SubscriptionFilter<R> + Send + Sync>,
    //     broadcast_tx: SubscriptionBroadcastTx<R>,
    //     respond_to: SubscriptionResponseTx,
    // },
    // /// Unsubscribe from messages from a module.
    // /// The payload should be used for additional subscription details required
    // /// for a specific implementation. The response will be
    // /// received via a oneshot Sender with message of type R.
    // Unsubscribe {
    //     subscription_id: usize,
    //     respond_to: SubscriptionResponseTx,
    // }
}

/// Control messages for managing the state of an AsyncIBusNode.
enum AsyncIBusNodeControlMessage {
    /// Shutdown the node.
    Shutdown { respond_to: oneshot::Sender<bool> },
    /// A specific request to the Node to change its opmode.
    /// The success of this request is indicated in the response.
    /// If successful, the response will match the request.
    /// If not, the response will indicate the OpMode to which the
    /// Node fell back.
    WriteOpMode {
        opmode: AsyncIBusOpMode,
        respond_to: oneshot::Sender<AsyncIBusOpMode>,
    },
    /// A specific request to poll the current state of the Node.
    ReadOpMode {
        respond_to: oneshot::Sender<AsyncIBusOpMode>,
    },
}

/// Common interface for Servelets.
#[async_trait]
pub trait AsyncIBusNode<S, R>: Send + Sync
where
    S: Clone + Send + 'static,
    R: Clone + Send + 'static,
{
    /// Returns the ID of this node.
    fn node_id(&self) -> String;

    /// A request to change the OpMode of this Node was received.
    /// A response is expected. If successful, respond with the matching opmode code. If not successful,
    /// respond with the opmode to which this Node fell back.
    /// # Parameters
    /// - `opmode`: The requested opmode.
    /// - `respond_to`: The oneshot channel to send the response with the resuling opmode.
    async fn request_change_opmode(
        &mut self,
        opmode: AsyncIBusOpMode,
        respond_to: oneshot::Sender<AsyncIBusOpMode>,
    );

    /// A synchronous version of request_change_opmode to allow for integration with Actix and other traits/actors
    /// with synchronous callbacks.
    // If successful, returns Ok(). If not successful,returns anyhow::Error.
    /// # Parameters
    /// - `opmode`: The requested opmode.
    fn request_change_opmode_sync(&mut self, opmode : AsyncIBusOpMode) -> Result<(), anyhow::Error>;

    /// A request to read the current OpMode of this node was recevied.
    /// A response is expected.
    /// # Parameters
    /// - `respond_to`: The oneshot channel to send the response with the resuling opmode.
    async fn request_read_opmode(&self, respond_to: oneshot::Sender<AsyncIBusOpMode>);


    /// A synchronous version of request_read_opmode.
    fn request_read_opmode_sync(&self) -> AsyncIBusOpMode;

    /// A command request was received. A response is expected.
    ///
    /// # Parameters
    /// - `msg`: The message payload of type S.
    /// - `respond_to`: The oneshot channel to send the response of type R.
    async fn request_received(&mut self, msg: S, respond_to: oneshot::Sender<R>);

    // /// A request to subscribe to a topic or events was received.
    // ///
    // /// # Parameters
    // /// - `msg`: The message payload of type S.
    // /// - `broadcast_tx`: The broadcast channel to send messages to the subscriber.
    // /// - `respond_to`: The oneshot channel to send the response of type R.
    // async fn subscribe_request_received(
    //     &mut self,
    //     msg: S,
    //     broadcast_tx: SubscriptionBroadcastTx<R>,
    //     respond_to: SubscriptionResponseTx,
    // );

    // /// A request to subscribe to a topic or events with a filter was received.
    // ///
    // /// # Parameters
    // /// - `msg`: The message payload of type S.
    // /// - `filter`: The filter to apply to the subscription.
    // /// - `broadcast_tx`: The broadcast channel to send messages to the subscriber.
    // /// - `respond_to`: The oneshot channel to send the response of type R.
    // async fn subscribe_with_filter_request_received(
    //     &mut self,
    //     msg: S,
    //     filter: Box<dyn SubscriptionFilter<S> + Send + Sync>,
    //     broadcast_tx: SubscriptionBroadcastTx<R>,
    //     respond_to: SubscriptionResponseTx,
    // );

    // /// A request to unsubscribe from a topic or events.
    // ///
    // /// # Parameters
    // /// - `msg`: The message payload of type S.
    // /// - `respond_to`: The oneshot channel to send the response of type R.
    // async fn unsubscribe_request_received(&mut self, msg: S, respond_to: oneshot::Sender<R>);

    /// A broadcasted message was received. This is an unsolicited message
    /// that does not expect a response.
    ///
    /// # Parameters
    /// - `msg`: The message payload of type R.
    async fn broadcast_received(&mut self, msg: R);

    // /// Send a message to a target node.
    // ///
    // /// # Parameters
    // /// - `target_node_id`: The ID of the target node.
    // /// - `registry_handle`: The registry handle to use for sending the message.
    // /// - `msg`: The message payload of type S.
    // ///
    // /// # Returns
    // /// - `Result<R, anyhow::Error>`: The response of type R or an error if the message could not be sent.
    // async fn send_message(
    //     &mut self,
    //     target_node_id: &str,
    //     registry_handle: AsyncNodeRegistryHandle<S, R>,
    //     msg: S,
    // ) -> Result<R, anyhow::Error> {
    //     let (send, recv) = oneshot::channel();
    //     let ibus_msg = IBusMessage::Request {
    //         payload: msg,
    //         respond_to: send,
    //     };
    //     match registry_handle.send_message(target_node_id, ibus_msg).await {
    //         Err(err) => {
    //             return Err(anyhow!(
    //                 "Failed to send message to {} with err: {}",
    //                 target_node_id,
    //                 err
    //             ));
    //         }
    //         _ => match recv.await {
    //             Ok(res) => {
    //                 return Ok(res);
    //             }
    //             Err(err) => {
    //                 return Err(anyhow!("Error receiving response to message: {}", err));
    //             }
    //         },
    //     }
    // }
}
/// Runs the IBus node, processing incoming messages and control messages.
///
/// # Parameters
/// - `node`: The IBus node to run.
/// - `ibus_rx`: The receiver for IBus messages.
/// - `control_rx`: The receiver for control messages.
///
///
///
///
///
/// Runs the IBus node, processing incoming messages and control messages.
///
/// # Parameters
/// - `node`: The IBus node to run.
/// - `ibus_rx`: The receiver for IBus messages.
/// - `control_rx`: The receiver for control messages.

async fn run_ibus_node<S, R>(
    mut node: Box<dyn AsyncIBusNode<S, R> + Send>,
    mut ibus_rx: mpsc::Receiver<IBusMessage<S, R>>,
    mut control_rx: mpsc::Receiver<AsyncIBusNodeControlMessage>,
) where
    S: Clone + Send + Sync + 'static,
    R: Clone + Send + Sync + 'static,
{
    let fin: Option<oneshot::Sender<bool>>;

    loop {
        tokio::select! {
            msg = ibus_rx.recv() => {
                if let Some(m) = msg {
                    match m {
                        IBusMessage::Request{payload, respond_to} => {
                            node.request_received(payload, respond_to).await;
                        },
                        IBusMessage::Broadcast{payload} => {

                            // alert the parent that we've received a broadcast
                            node.broadcast_received(payload).await;

                        },
                        // IBusMessage::Subscribe{topic, broadcast_tx, respond_to} => {

                        //     // Forward the subscription request to the SubscriptionManager
                        //     subscription_manager.subscribe(topic, broadcast_tx, respond_to).await;


                        //     // let subscription_manager_clone = subscription_manager.clone();
                        //     // tokio::spawn(async move {
                        //     //     subscription_manager_clone.lock().await.subscribe(
                        //     //         topic,
                        //     //         broadcast_tx,
                        //     //         respond_to,
                        //     //     ).await;
                        //     // });
                        // },
                        // IBusMessage::Unsubscribe{subscription_id, respond_to} => {

                        //     // Forward the unsubscribe request to the SubscriptionManager
                        //     subscription_manager.unsubscribe(subscription_id, respond_to).await;

                        //     // let subscription_manager_clone = subscription_manager.clone();
                        //     // tokio::spawn(async move {
                        //     //     subscription_manager_clone.lock().await.unsubscribe(
                        //     //         subscription_id,
                        //     //         respond_to,
                        //     //     ).await;
                        //     // });
                        // },
                        // IBusMessage::SubscribeWithFilter { topic, filter, broadcast_tx, respond_to } => {
                        //     // Forward the subscription request with filter to the SubscriptionManager
                        //     subscription_manager.subscribe_with_filter(topic, filter, broadcast_tx, respond_to).await;

                        //     // let subscription_manager = subscription_manager.clone();
                        //     // tokio::spawn(async move {
                        //     //     subscription_manager.lock().await.subscribe_with_filter(
                        //     //         topic,
                        //     //         filter,
                        //     //         broadcast_tx,
                        //     //         respond_to,
                        //     //     ).await;
                        //     // });
                        // }
                    }
                }
            },
            ctrl = control_rx.recv() => {
                if let Some(m) = ctrl {
                    match m {
                        AsyncIBusNodeControlMessage::Shutdown {respond_to} => {
                            fin = Some(respond_to);
                            break;
                        },
                        AsyncIBusNodeControlMessage::WriteOpMode { opmode, respond_to } => {
                            node.request_change_opmode(opmode, respond_to).await;
                        },
                        AsyncIBusNodeControlMessage::ReadOpMode { respond_to } => {
                            node.request_read_opmode(respond_to).await;
                        }
                    }
                }
            }
        }
    }

    debug!("AsyncIBusNode pipeline shutting down...");

    if let Some(sender) = fin {
        let _ = sender.send(true);
    }
}

/// Handle for controlling an AsyncIBusNode.
pub struct AsyncIBusNodeHandle<S, R>
where
    S: Clone + Send + Sync + 'static,
    R: Clone + Send + Sync + 'static,
{
    control_tx: mpsc::Sender<AsyncIBusNodeControlMessage>,
    registry_handle: AsyncNodeRegistryHandle<S, R>,
    node_id: String,

    /// PhantomData to hold the types
    _marker: PhantomData<(S, R)>,
}

impl<S, R> AsyncIBusNodeHandle<S, R>
where
    S: Clone + Send + Sync + 'static,
    R: Clone + Send + Sync + 'static,
{
    /// Creates a new AsyncIBusNodeHandle.
    ///
    /// # Parameters
    /// - `node`: The IBus node to handle.
    /// - `ibus_rx`: The receiver for IBus messages.
    /// - `registry_handle`: The registry handle for managing node communication.
    ///
    /// # Returns
    /// - `Self`: The created AsyncIBusNodeHandle.
    pub fn new(
        node: Box<dyn AsyncIBusNode<S, R> + Send>,
        ibus_rx: mpsc::Receiver<IBusMessage<S, R>>,
        registry_handle: AsyncNodeRegistryHandle<S, R>,
    ) -> Self {
        // Create the communication channels.
        let (ctrl_tx, ctrl_rx) = mpsc::channel(8);

        // Create an instance of this handle. This is the ibus interface to the derived instance
        // of the AsyncIBusNode trait.
        let ret = Self {
            control_tx: ctrl_tx,
            registry_handle,
            node_id: node.node_id(),
            _marker: PhantomData,
        };

        // Start processing the ibus node in a separate task.
        tokio::spawn(run_ibus_node(node, ibus_rx, ctrl_rx));

        return ret;
    }

    /// Returns the ID of the node.
    ///
    /// # Returns
    /// - `&str`: The ID of the node.
    pub fn node_id(&self) -> &str {
        return self.node_id.as_str();
    }

    /// Sends a request message to the node and waits for a response.
    ///
    /// # Parameters
    /// - `msg`: The message payload of type S.
    ///
    /// # Returns
    /// - `Result<R, anyhow::Error>`: The response of type R or an error if the message could not be sent.
    pub async fn request(&self, msg: S) -> Result<R, anyhow::Error> {
        let (send, recv) = oneshot::channel();
        let ibus_msg = IBusMessage::Request {
            payload: msg,
            respond_to: send,
        };
        match self
            .registry_handle
            .send_message(&self.node_id, ibus_msg)
            .await
        {
            Ok(_) => match recv.await {
                Ok(res) => {
                    return Ok(res);
                }
                Err(err) => {
                    return Err(anyhow!("Error receiving response to message: {}", err));
                }
            },
            Err(err) => {
                return Err(anyhow!(
                    "Failed to send message to node {} with err: {}",
                    self.node_id,
                    err
                ));
            }
        }
    }

    /// Sends a broadcast message to the node.
    ///
    /// # Parameters
    /// - `payload`: The message payload of type R.
    ///
    /// # Returns
    /// - `Result<(), anyhow::Error>`: An empty result or an error if the message could not be sent.
    pub async fn broadcast(&self, payload: R) -> Result<(), anyhow::Error> {
        let ibus_msg = IBusMessage::Broadcast { payload: payload };
        match self
            .registry_handle
            .send_message(&self.node_id, ibus_msg)
            .await
        {
            Ok(_) => {
                return Ok(());
            }
            Err(err) => {
                return Err(anyhow!(
                    "Failed to send message to node {} with err: {}",
                    self.node_id,
                    err
                ));
            }
        }
    }

    /// Sends a request to change the operational mode of the node and waits for a response.
    ///
    /// # Parameters
    /// - `opmode`: The requested operational mode.
    ///
    /// # Returns
    /// - `Result<AsyncIBusOpMode, anyhow::Error>`: The resulting operational mode or an error if the request could not be sent.
    pub async fn request_change_opmode(
        &self,
        opmode: AsyncIBusOpMode,
    ) -> Result<AsyncIBusOpMode, anyhow::Error> {
        let (send, recv) = oneshot::channel();
        let ctrl_msg = AsyncIBusNodeControlMessage::WriteOpMode {
            opmode: opmode,
            respond_to: send,
        };
        match self.control_tx.send(ctrl_msg).await {
            Ok(_) => match recv.await {
                Ok(res) => {
                    return Ok(res);
                }
                Err(err) => {
                    return Err(anyhow!("Error receiving response to message: {}", err));
                }
            },
            Err(err) => {
                return Err(anyhow!(
                    "Failed to send message to node {} with err: {}",
                    self.node_id,
                    err
                ));
            }
        }
    }

    /// Sends a request to read the current operational mode of the node and waits for a response.
    ///
    /// # Returns
    /// - `Result<AsyncIBusOpMode, anyhow::Error>`: The current operational mode or an error if the request could not be sent.
    pub async fn request_read_opmode(&self) -> Result<AsyncIBusOpMode, anyhow::Error> {
        let (send, recv) = oneshot::channel();
        let ctrl_msg = AsyncIBusNodeControlMessage::ReadOpMode { respond_to: send };
        match self.control_tx.send(ctrl_msg).await {
            Ok(_) => match recv.await {
                Ok(res) => {
                    return Ok(res);
                }
                Err(err) => {
                    return Err(anyhow!("Error receiving response to message: {}", err));
                }
            },
            Err(err) => {
                return Err(anyhow!(
                    "Failed to send message to node {} with err: {}",
                    self.node_id,
                    err
                ));
            }
        }
    }

    /// Internal control message.
    /// Requests the node to shutdown. This will stop all processing of messages to the node and rendering
    /// it unusable. This will lead to it destructing itself.
    /// # Returns
    /// - `Result<(), anyhow::Error>`: Ok or error if the request could not be sent.
    pub async fn shutdown(&self) -> Result<(), anyhow::Error> {
        let (send, recv) = oneshot::channel();
        let ctrl_msg = AsyncIBusNodeControlMessage::Shutdown { respond_to: send };
        match self.control_tx.send(ctrl_msg).await {
            Ok(_) => match recv.await {
                Ok(_) => {
                    return Ok(());
                }
                Err(err) => {
                    return Err(anyhow!("Error receiving response to message: {}", err));
                }
            },
            Err(err) => {
                return Err(anyhow!(
                    "Failed to send message to node {} with err: {}",
                    self.node_id,
                    err
                ));
            }
        }
    }
}

/// This AsyncIBusNode Factory helps streamline and standardize the creation of
/// structs implementing AsyncIBusNode trait.
pub trait AsyncIBusNodeFactory<S, R, CTX> {
    /// Create a new instance of the AsyncIBusNode. Returns a handle to the created instance if
    /// successful, an error if the node fails to register with the provided node registry.
    ///
    /// # Parameters
    /// - `node_id`: The ID of the node to create.
    /// - `registry_handle`: The registry handle for managing node communication.
    /// - `configuration`: The configuration for the node.
    ///
    /// # Returns
    /// - `Result<AsyncIBusNodeHandle<S, R>, anyhow::Error>`: The created AsyncIBusNodeHandle or an error if the node fails to register.
    fn create(
        node_id: &str,
        registry_handle: AsyncNodeRegistryHandle<S, R>,
        configuration: &serde_json::Value,
        context : CTX
    ) -> Result<AsyncIBusNodeHandle<S, R>, anyhow::Error>
    where
        Self: Sized + AsyncIBusNode<S, R> + 'static,
        S: Clone + Send + Sync + 'static,
        R: Clone + Send + Sync + 'static,
    {
        let (ibus_tx, ibus_rx) = mpsc::channel(64);

        let handle;
        match Self::init_node(node_id.to_string(), registry_handle.clone(), configuration, context) {
            Ok(ret) => {
                handle = AsyncIBusNodeHandle::<S, R>::new(ret, ibus_rx, registry_handle.clone());
            }
            Err(err) => {
                return Err(anyhow!("Error initializing node {} : {}", node_id, err));
            }
        }

        // Block on running the ibus node (non-async function waiting for async to finish)
        let res = tokio::task::block_in_place(|| {
            let runtime = tokio::runtime::Handle::current();
            return runtime.block_on(registry_handle.register_node(node_id, ibus_tx));
        });

        match res {
            Ok(_) => return Ok(handle),
            Err(err) => {
                return Err(anyhow!(
                    "Failed to register new node with registry: {}",
                    err
                ));
            }
        }
    }

    

    /// To be implemented by the AsyncIBusNode. The instance should initialize itself, storing
    /// the node_id and registry_handle, both of which it will need to participate in the bus.
    ///
    /// # Parameters
    /// - `node_id`: The ID of the node to initialize.
    /// - `registry_handle`: The registry handle for managing node communication.
    /// - `configuration`: The configuration for the node.
    ///
    /// # Returns
    /// - `Self`: The initialized node.
    fn init_node(
        node_id: String,
        registry_handle: AsyncNodeRegistryHandle<S, R>,
        configuration: &serde_json::Value,
        context : CTX
    ) -> Result<Box<Self>, anyhow::Error>;
}