ProjectNautilus/ground/src/command.rs

use anyhow::bail;
use api::client::Client;
use api::client::command::CommandRegistry;
use api::macros::IntoCommandDefinition;
use chrono::{DateTime, TimeDelta, Utc};
use log::{error, trace};
use nautilus_common::command::set_pin::SetPin;
use nautilus_common::command::set_rcs::SetRcs;
use nautilus_common::command::valid_priority_command::ValidPriorityCommand;
use nautilus_common::command::{Command, OwnedCommandHeader};
use nautilus_common::math::Vector;
use nautilus_common::udp::tokio::AsyncUdpSocketExt;
use std::fmt::Debug;
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::net::UdpSocket;
use tokio::select;
use tokio::sync::RwLock;
use tokio::sync::mpsc::{Sender, channel};
use tokio_util::sync::CancellationToken;

const MAX_DATETIME: DateTime<Utc> = DateTime::from_timestamp_nanos(i64::MAX);

pub struct CommandHandler<'a> {
    cmd: CommandRegistry,
    flight_addr: &'a RwLock<SocketAddr>,
    udp: &'a UdpSocket,
    cancel: CancellationToken,
}

#[derive(IntoCommandDefinition)]
struct SetPinCommand {
    index: u8,
    state: bool,
}
impl From<SetPinCommand> for SetPin {
    fn from(value: SetPinCommand) -> Self {
        Self {
            pin: value.index,
            value: value.state,
        }
    }
}

#[derive(IntoCommandDefinition)]
struct RcsCommand {
    translate_x: f64,
    translate_y: f64,
    translate_z: f64,
    rotate_x: f64,
    rotate_y: f64,
    rotate_z: f64,
}
impl From<RcsCommand> for SetRcs {
    fn from(value: RcsCommand) -> Self {
        Self {
            translation: Vector::new(value.translate_x, value.translate_y, value.translate_z),
            rotation: Vector::new(value.rotate_x, value.rotate_y, value.rotate_z),
        }
    }
}

impl<'a> CommandHandler<'a> {
    pub fn new(
        client: Arc<Client>,
        flight_addr: &'a RwLock<SocketAddr>,
        udp: &'a UdpSocket,
        cancel: CancellationToken,
    ) -> Self {
        Self {
            cmd: CommandRegistry::new(client),
            flight_addr,
            udp,
            cancel,
        }
    }

    pub fn run(self) -> anyhow::Result<()> {
        let runtime = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()?;

        runtime.block_on(async move {
            let (outgoing_commands_tx, mut outgoing_commands_rx) = channel::<OwnedCommandHeader>(16);

            let mut commands = ["a", "b"].iter()
                .map(|bank| {
                    let command_name = format!("/mcp23017{bank}/set");
                    let outgoing_commands_tx = outgoing_commands_tx.clone();
                    self.cmd.register_handler(
                        format!("switch.{bank}.set"),
                        move |_, cmd: SetPinCommand| -> anyhow::Result<_> {
                            trace!("Setting Switch {bank} {}", cmd.index);

                            if cmd.index >= 16 {
                                bail!("Invalid Pin Number: {}", cmd.index)
                            }

                            outgoing_commands_tx.try_send_command(
                                &command_name,
                                &ValidPriorityCommand {
                                    inner: SetPin::from(cmd),
                                    valid_until: MAX_DATETIME,
                                    priority: 0,
                                }
                            )?;

                            Ok("Command Executed Successfully".to_string())
                        }
                    )
                })
                .collect::<Vec<_>>();

            {
                let outgoing_commands_tx = outgoing_commands_tx.clone();
                commands.push(self.cmd.register_handler("shutdown", move |_, ()| -> anyhow::Result<_> {
                    trace!("Shutting Down Flight");

                    outgoing_commands_tx.try_send_command("/shutdown", &())?;

                    Ok("Command Executed Successfully".to_string())
                }));
            }

            {
                let outgoing_commands_tx = outgoing_commands_tx.clone();
                commands.push(self.cmd.register_handler("rcs.set", move |header, cmd: RcsCommand| -> anyhow::Result<_> {
                    trace!("Sending Rcs Set Command");

                    outgoing_commands_tx.try_send_command(
                        "/rcs/set",
                        &ValidPriorityCommand {
                            inner: SetRcs::from(cmd),
                            valid_until: header.timestamp + TimeDelta::seconds(5),
                            priority: 0,
                        }
                    )?;

                    Ok("Command Executed Successfully".to_string())
                }));
            }

            // We no longer need this
            drop(outgoing_commands_tx);

            let mut buffer = [0u8; 512];
            while !self.cancel.is_cancelled() {
                select! {
                    () = self.cancel.cancelled() => { break; }
                    outgoing = outgoing_commands_rx.recv() => {
                        match outgoing {
                            None => break,
                            Some(outgoing) => {
                                match self.udp.send_postcard(&outgoing, &mut buffer, *self.flight_addr.read().await).await {
                                    Ok(_) => {},
                                    Err(err) => error!("Failed to Send Outgoing {err}"),
                                }
                            }
                        }
                    }
                }
            }

            // Explicit Drops
            drop(commands);
            drop(self);
        });

        Ok(())
    }
}

trait SenderExt {
    fn try_send_command<T: Command + Debug>(&self, name: &str, data: &T) -> anyhow::Result<()>;
}

impl SenderExt for Sender<OwnedCommandHeader> {
    fn try_send_command<T: Command + Debug>(&self, name: &str, data: &T) -> anyhow::Result<()> {
        trace!("{data:?}");
        let inner_buffer = postcard::to_allocvec(data)?;
        self.try_send(OwnedCommandHeader {
            name: name.to_string(),
            data: inner_buffer,
        })?;
        Ok(())
    }
}