//
// Copyright (C) 2024 Automated Design Corp.. All Rights Reserved.
// Created Date: 2024-09-03 11:24:02
// -----
// Last Modified: 2024-09-05 15:09:08
// -----
//
//

//! # MechActor: The MechUtil Actor
//!
//! The `MechActor` is an opinionated actor model for managing asynchronous tasks and routing commands
//! in a decoupled and efficient way, primarily using `tokio` channels. This actor provides both synchronous
//! and asynchronous tick handling, as well as routing functionality for commands.
//!
//! ## Motivation
//!
//! In mechatronics and automation, often the same resource (e.g., a connection to a controller) needs
//! to be shared across multiple asynchronous tasks performing unrelated functions. The `MechActor` solves
//! the following problems:
//! - Avoids repetitive code for spawning async tasks.
//! - Routes commands easily throughout a decoupled application.
//! - Works well in the context of a structured implementation without the complexities of lifetimes.
//!
//! ## Design Inspiration
//!
//! The actor's design is inspired by Alice Ryhl’s blog post [Actors with Tokio](https://ryhl.io/blog/actors-with-tokio/).
//! It uses `tokio` channels to avoid lifetime issues common with other actor models such as Actix.
//!
//! ## Features
//! - Regularly scheduled tick events at a specified interval, regardless of active usage.
//! - Support for both synchronous and asynchronous callbacks.
//! - Simple API for registering command routes.
//!
//! ## Usage
//!
//! ```ignore
//! use tokio::time::Duration;
//! use mechutil::MechActor;
//!
//! // Example of using the Actor in a struct
//! struct ActorStructExample {
//!     actor : MechActor
//! }
//!
//! impl ActorStructExample {
//!     pub fn new() -> Self {
//!         Self {
//!             actor : MechActor::new(std::time::Duration::from_millis(1000))            
//!         }
//!     }
//!
//!
//!     pub async fn start(&mut self) {
//!         let _ = self.actor.register_tick( move || {
//!             println!("actor 2 tick 1 ");             
//!         }).await;
//!
//!
//!         let _ = self.actor.register_tick( move || {
//!             println!("actor 2 tick 2 ");             
//!         }).await;
//!
//!
//!         if let Err(err) = self.actor.register_async_tick( move || {
//!             Box::pin (async move {
//!                 println!("async actor 2 TICKING ASYNC, BABY!");
//!             })
//!         }).await {
//!             println!("Failed to register async tick: {}", err);
//!         }
//!
//!         println!("SpawnTest registered tick callback");
//!     }
//!
//!     pub async fn stop(&mut self) {
//!         let _ = self.actor.shutdown().await;
//!         println!("SpawnTest shutdown.");
//!     }
//! }
//!
//!
//! #[tokio::main]
//! async fn main() {
//!     let actor = MechActor::new(Duration::from_secs(1));
//!
//!     // Register a synchronous tick callback
//!     let _ = actor.register_tick(|| {
//!         println!("Sync tick executed!");
//!     }).await;
//!
//!     // Register an asynchronous tick callback
//!     let _ = actor.register_async_tick(|| {
//!         Box::pin(async {
//!             println!("Async tick executed!");
//!         })
//!     }).await;
//!
//!     // Register a command route
//!     let _ = actor.register_route("example", |cmd| {
//!         println!("Processing command: {:?}", cmd);
//!         cmd.clone() // Return a processed command
//!     }).await;
//!
//!     // Simulate shutdown after some time
//!     tokio::time::sleep(Duration::from_secs(5)).await;
//!     let _ = actor.shutdown().await;
//! }
//! ```

use anyhow::anyhow;
use std::{collections::HashMap, pin::Pin, time::Duration};
use tokio::sync::{mpsc, oneshot};

use std::future::Future;

use crate::command::CommandMessage;

type BoxedCallback = Box<dyn FnMut(&CommandMessage) -> CommandMessage + Send + 'static>;
type BoxedTickCallback = Box<dyn FnMut() -> () + Send + 'static>;

type BoxedAsyncCallback = Box<
    dyn FnMut(&CommandMessage) -> Pin<Box<dyn Future<Output = CommandMessage> + Send>>
        + Send
        + 'static,
>;
type BoxedAsyncTickCallback =
    Box<dyn FnMut() -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + 'static>;

/// Messages that can be sent to the `MechActor`.
enum ActorMessage {
    RegisterRoute {
        respond_to: oneshot::Sender<u32>,
        route: String,
        cb: BoxedCallback,
    },
    ProcessRoute {
        respond_to: oneshot::Sender<u32>,
        route: String,
        arg: CommandMessage,
    },
    RegisterAsyncRoute {
        respond_to: oneshot::Sender<u32>,
        route: String,
        cb: BoxedAsyncCallback,
    },
    RegisterTick {
        respond_to: oneshot::Sender<u32>,
        cb: BoxedTickCallback,
    },
    RegisterAsyncTick {
        respond_to: oneshot::Sender<u32>,
        cb: BoxedAsyncTickCallback,
    },
    Shutdown,
}

/// Internal actor processes the messages and contains the registered callbacks.
struct MechActorInternal {
    receiver: mpsc::Receiver<ActorMessage>,

    next_id: u32,
    routes: HashMap<String, BoxedCallback>,
    tick_routes: HashMap<u32, BoxedTickCallback>,
    tick_uid: u32,
    async_routes: HashMap<String, BoxedAsyncCallback>,
    async_tick_routes: HashMap<u32, BoxedAsyncTickCallback>,
}

// Implement the internal actor
impl MechActorInternal {
    fn new(receiver: mpsc::Receiver<ActorMessage>) -> Self {
        MechActorInternal {
            receiver,
            next_id: 0,
            routes: HashMap::new(),
            tick_routes: HashMap::new(),
            tick_uid: 0,
            async_routes: HashMap::new(),
            async_tick_routes: HashMap::new(),
        }
    }

    fn handle_message(&mut self, msg: ActorMessage) {
        match msg {
            ActorMessage::RegisterRoute {
                respond_to,
                route,
                cb,
            } => {
                self.routes.insert(route, cb);
                let _ = respond_to.send(self.next_id);
            }
            ActorMessage::RegisterTick { respond_to, cb } => {
                self.tick_uid += 1;
                self.tick_routes.insert(self.tick_uid, cb);
                let _ = respond_to.send(self.tick_uid);
            }
            ActorMessage::RegisterAsyncRoute {
                respond_to,
                route,
                cb,
            } => {
                self.async_routes.insert(route, cb);
                let _ = respond_to.send(self.next_id);
            }
            ActorMessage::RegisterAsyncTick { respond_to, cb } => {
                self.tick_uid += 1;
                self.async_tick_routes.insert(self.tick_uid, cb);
                let _ = respond_to.send(self.tick_uid);
            }
            _ => {
                // Unknown message. Ignore.
            }
        }
    }

    async fn tick(&mut self) {
        for (_, cb) in &mut self.tick_routes {
            cb();
        }

        for (_, cb) in &mut self.async_tick_routes {
            cb().await;
        }
    }
}

/// Execute the internal actor. Runs the tick interval and pipes messages into the
/// internal actor. This is also where routes are processed.
async fn run_my_actor(mut actor: MechActorInternal, interval: Duration) {
    let mut ticker = tokio::time::interval(interval);

    loop {
        tokio::select! {
            _ = ticker.tick() => {
                actor.tick().await;
            },
            msg = actor.receiver.recv() => {

                if let Some(m) = msg {
                    match m {
                        ActorMessage::Shutdown => {
                            break;
                        },
                        ActorMessage::ProcessRoute { respond_to : _, route, arg} => {
                            let arg_copy = arg.clone();
                            if let Some(cb) = actor.async_routes.get_mut(&route) {
                                cb(&arg_copy).await;
                            }
                            if let Some(cb) = actor.routes.get_mut(&route) {
                                cb(&arg_copy);
                            }
                        },
                        _ => {
                            actor.handle_message(m);
                        }
                    }
                }

            }
        }
    }
}

/// The public API for interacting with a `MechActor`.
#[derive(Clone)]
pub struct MechActor {
    sender: mpsc::Sender<ActorMessage>,
}

impl MechActor {
    /// Creates a new `MechActor` that ticks at the specified `interval`.
    ///
    /// # Arguments
    ///
    /// * `interval` - The time duration between each tick.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// use mechutil::MechActor;
    /// use std::time::Duration;
    ///
    /// let actor = MechActor::new(Duration::from_secs(1));
    /// ```    
    pub fn new(interval: Duration) -> Self {
        let (sender, receiver) = mpsc::channel(8);
        let actor = MechActorInternal::new(receiver);
        tokio::spawn(run_my_actor(actor, interval));

        Self { sender: sender }
    }

    /// Registers a synchronous callback to be invoked on a matching route.
    ///
    /// # Arguments
    ///
    /// * `cb` - The callback function to be invoked on each tick.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// let actor = MechActor::new(Duration::from_secs(1));
    /// let _ = actor.register_route( "MyRoute".to_string(), || {
    ///     println!("Tick happened!");
    /// }).await;
    /// ```        
    pub async fn register_route<C>(&self, route: &str, cb: C) -> Result<(), anyhow::Error>
    where
        C: Fn(&CommandMessage) -> CommandMessage + Send + 'static,
    {
        let (send, recv) = oneshot::channel();
        let msg = ActorMessage::RegisterRoute {
            respond_to: send,
            route: route.to_string(),
            cb: Box::new(cb),
        };

        // Ignore send errors. If this send fails, so does the
        // recv.await below. There's no reason to check for the
        // same failure twice.
        let _ = self.sender.send(msg).await;
        match recv.await {
            Ok(_) => Ok(()),
            Err(err) => return Err(anyhow!("Error registering route: {}", err)),
        }
    }

    /// Registers a synchronous callback to be invoked on each tick.
    ///
    /// # Arguments
    ///
    /// * `cb` - The callback function to be invoked on each tick.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// let actor = MechActor::new(Duration::from_secs(1));
    /// let _ = actor.register_tick(|| {
    ///     println!("Tick happened!");
    /// }).await;
    /// ```    
    pub async fn register_tick<C>(&self, mut cb: C) -> Result<(), anyhow::Error>
    where
        C: FnMut() -> () + Send + 'static,
    {
        let (send, recv) = oneshot::channel();
        let msg = ActorMessage::RegisterTick {
            respond_to: send,
            cb: Box::new(move || cb()),
        };

        // Ignore send errors. If this send fails, so does the
        // recv.await below. There's no reason to check for the
        // same failure twice.
        let _ = self.sender.send(msg).await;
        match recv.await {
            Ok(_) => Ok(()),
            Err(err) => return Err(anyhow!("Error registering tick: {}", err)),
        }
    }

    /// Registers an asynchronous callback to be invoked on a matching route.
    ///
    /// # Arguments
    ///
    /// * `route` - String identifying the route.
    /// * `cb` - The async callback function to be invoked on each tick.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// let actor = MechActor::new(Duration::from_secs(1));
    /// let _ = actor.register_async_route( "MyRoute".to_string(), || {
    ///     Box::pin(async {
    ///         println!("Async tick happened!");
    ///     })
    /// }).await;
    /// ```        
    pub async fn register_async_route<C, Fut>(
        &self,
        route: &str,
        cb: C,
    ) -> Result<(), anyhow::Error>
    where
        C: Fn(&CommandMessage) -> Fut + Send + 'static,
        Fut: Future<Output = CommandMessage> + Send + 'static,
    {
        let (send, recv) = oneshot::channel();
        let msg = ActorMessage::RegisterAsyncRoute {
            respond_to: send,
            route: route.to_string(),
            cb: Box::new(move |cmd| Box::pin(cb(cmd))),
        };

        let _ = self.sender.send(msg).await;
        match recv.await {
            Ok(_) => Ok(()),
            Err(err) => return Err(anyhow!("Error registering route: {}", err)),
        }
    }

    /// Registers an asynchronous callback to be invoked on each tick.
    ///
    /// # Arguments
    ///
    /// * `cb` - The async callback function to be invoked on each tick.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// let actor = MechActor::new(Duration::from_secs(1));
    /// let _ = actor.register_async_tick(|| {
    ///     Box::pin(async {
    ///         println!("Async tick happened!");
    ///     })
    /// }).await;
    /// ```    
    pub async fn register_async_tick<C, Fut>(&self, mut cb: C) -> Result<(), anyhow::Error>
    where
        C: FnMut() -> Fut + Send + 'static,
        Fut: Future<Output = ()> + Send + 'static,
    {
        let (send, recv) = oneshot::channel();
        let msg = ActorMessage::RegisterAsyncTick {
            respond_to: send,
            cb: Box::new(move || Box::pin(cb())),
        };

        let _ = self.sender.send(msg).await;
        match recv.await {
            Ok(_) => Ok(()),
            Err(err) => return Err(anyhow!("Error registering tick: {}", err)),
        }
    }

    /// Shuts down the `MechActor`.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// let _ = actor.shutdown().await;
    /// ```    
    pub async fn shutdown(&self) -> Result<(), anyhow::Error> {
        let (_, recv) = oneshot::channel::<u32>();
        let msg = ActorMessage::Shutdown;

        // Ignore send errors. If this send fails, so does the
        // recv.await below. There's no reason to check for the
        // same failure twice.
        let _ = self.sender.send(msg).await;

        match recv.await {
            Ok(_) => Ok(()),
            Err(err) => return Err(anyhow!("Error shutting down Actor: {}", err)),
        }
    }
}
