main.rs

use std::{env, future::Future, path::PathBuf, sync::Arc, time::Duration};

use clap::Parser;
use color_eyre::{
    Result,
    eyre::{Context as _, ContextCompat, bail, eyre},
};
use repository::{Repository, team::Team};
use ros_z::prelude::*;
use tokio::task::JoinSet;
use tracing_subscriber::EnvFilter;

const RUNTIME_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(2);

#[derive(Debug, Parser)]
struct Args {
    #[arg(long)]
    location: String,
    #[arg(long, default_value = "parameters/ros_z")]
    parameter_root: PathBuf,
    #[arg(long)]
    router: Option<String>,
    #[arg(long)]
    log_path: Option<PathBuf>,
}

struct RunningStack {
    join_set: JoinSet<Result<()>>,
}

fn main() -> Result<()> {
    color_eyre::install()?;
    tracing_subscriber::fmt()
        .with_env_filter(EnvFilter::from_default_env())
        .init();

    run_with_shutdown_timeout(run(), RUNTIME_SHUTDOWN_TIMEOUT)?
}

fn run_with_shutdown_timeout<F>(future: F, shutdown_timeout: Duration) -> Result<F::Output>
where
    F: Future,
{
    let runtime = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .wrap_err("failed to build Tokio runtime")?;
    let output = runtime.block_on(future);
    runtime.shutdown_timeout(shutdown_timeout);
    Ok(output)
}

async fn run() -> Result<()> {
    let args = Args::parse();

    let Some(hardware_id) = env::var_os("HARDWARE_ID") else {
        bail!("environment variable HARDWARE_ID not set");
    };
    let hardware_id = hardware_id
        .into_string()
        .ok()
        .wrap_err("id was not valid UTF-8")?;
    let robot_number = load_robot_number(&hardware_id).await?;
    let namespace = derive_namespace(&robot_number.to_string());

    let parameter_layers =
        derive_parameter_layers(&args.parameter_root, &args.location, &hardware_id);

    let mut builder = ContextBuilder::default()
        .with_namespace(&namespace)
        .with_parameter_layers(parameter_layers);

    builder = match args.router {
        Some(router) => builder.with_mode("client").with_router_endpoint(router)?,
        None => builder
            .with_mode("router")
            .disable_multicast_scouting()
            .with_connect_endpoints(std::iter::empty::<&str>())
            .with_listen_endpoints(["tcp/127.0.0.1:7447"]),
    };

    let ctx = Arc::new(builder.build().await?);
    let mut running = spawn_all(ctx.clone(), args.log_path).await?;

    let result = tokio::select! {
        result = monitor(&mut running.join_set) => result,
        _ = tokio::signal::ctrl_c() => {
            Ok(())
        }
    };

    running.join_set.abort_all();
    if result.is_ok() {
        ctx.shutdown()?;
    }
    result
}

fn derive_parameter_layers(
    parameter_root: &std::path::Path,
    location: &str,
    robot: &str,
) -> Vec<PathBuf> {
    vec![
        parameter_root.join("base"),
        parameter_root.join("location").join(location),
        parameter_root.join("robot").join(robot),
    ]
}

async fn load_robot_number(hardware_id: &str) -> Result<u8> {
    let repository =
        Repository::new(env::current_dir().wrap_err("failed to get current directory")?);
    let team = repository.read_team_configuration().await?;
    robot_number_for_hardware_id(&team, hardware_id)
}

fn robot_number_for_hardware_id(team: &Team, hardware_id: &str) -> Result<u8> {
    team.robots
        .iter()
        .find(|robot| robot.id == hardware_id)
        .map(|robot| robot.number)
        .ok_or_else(|| eyre!(r#"ID "{hardware_id}" not found in team.toml"#))
}

fn derive_namespace(robot: &str) -> String {
    if robot.starts_with('/') {
        robot.to_string()
    } else {
        format!("/{robot}")
    }
}

async fn spawn_all(ctx: Arc<Context>, log_path: Option<PathBuf>) -> Result<RunningStack> {
    let mut join_set = JoinSet::new();

    join_set.spawn(active_vision::run_boxed(ctx.clone()));
    join_set.spawn(ball_filter::run_boxed(ctx.clone()));
    join_set.spawn(ball_state_composer::run_boxed(ctx.clone()));
    join_set.spawn(behavior_node::run_boxed(ctx.clone()));
    join_set.spawn(booster_sdk_interface::run_boxed(ctx.clone()));
    join_set.spawn(button_event_bridge::run_boxed(ctx.clone()));
    join_set.spawn(button_event_handler::run_boxed(ctx.clone()));
    join_set.spawn(camera_matrix_calculator::run_boxed(ctx.clone()));
    join_set.spawn(detection::run_boxed(ctx.clone()));
    join_set.spawn(fake_odometry::run_boxed(ctx.clone()));
    join_set.spawn(fall_down_state_receiver::run_boxed(ctx.clone()));
    join_set.spawn(field_mark_association::run_boxed(ctx.clone()));
    join_set.spawn(game_controller_filter::run_boxed(ctx.clone()));
    join_set.spawn(game_controller_state_filter::run_boxed(ctx.clone()));
    join_set.spawn(global_parameter_provider::run_boxed(ctx.clone()));
    join_set.spawn(ground_provider::run_boxed(ctx.clone()));
    join_set.spawn(head_motion::run_boxed(ctx.clone()));
    join_set.spawn(image_receiver::run_boxed(ctx.clone()));
    join_set.spawn(kinematics_provider::run_boxed(ctx.clone()));
    join_set.spawn(led_handler::run_boxed(ctx.clone()));
    join_set.spawn(localization_2d::run_boxed(ctx.clone()));
    join_set.spawn(localization_3d::run_boxed(ctx.clone()));
    join_set.spawn(look_around::run_boxed(ctx.clone()));
    join_set.spawn(look_at::run_boxed(ctx.clone()));
    join_set.spawn(low_state_bridge::run_boxed(ctx.clone()));
    join_set.spawn(mcap_recorder::run_boxed(ctx.clone(), log_path));
    join_set.spawn(message_filter::run_boxed(ctx.clone()));
    join_set.spawn(message_handler::run_boxed(ctx.clone()));
    join_set.spawn(microphone_recorder::run_boxed(ctx.clone()));
    join_set.spawn(motor_commands_collector::run_boxed(ctx.clone()));
    join_set.spawn(obstacle_filter::run_boxed(ctx.clone()));
    join_set.spawn(odometer_bridge::run_boxed(ctx.clone()));
    join_set.spawn(odometry::run_boxed(ctx.clone()));
    join_set.spawn(player_states_receiver::run_boxed(ctx.clone()));
    join_set.spawn(primary_state_filter::run_boxed(ctx.clone()));
    join_set.spawn(visual_kick_ball_selector::run_boxed(ctx.clone()));
    join_set.spawn(rule_obstacle_composer::run_boxed(ctx.clone()));
    join_set.spawn(safe_pose_checker::run_boxed(ctx.clone()));
    join_set.spawn(search_suggestor::run_boxed(ctx.clone()));
    join_set.spawn(segment_filter::run_boxed(ctx.clone()));
    join_set.spawn(stereo_visual_odometry::run_boxed(ctx.clone()));
    join_set.spawn(support_foot_estimator::run_boxed(ctx.clone()));
    join_set.spawn(team_ball_receiver::run_boxed(ctx.clone()));
    join_set.spawn(time_to_reach_kick_position::run_boxed(ctx.clone()));
    join_set.spawn(trigger::run_boxed(ctx.clone()));
    join_set.spawn(whistle_detection::run_boxed(ctx.clone()));
    join_set.spawn(whistle_filter::run_boxed(ctx.clone()));
    join_set.spawn(world_state_composer::run_boxed(ctx.clone()));
    join_set.spawn(world_to_field_provider::run_boxed(ctx.clone()));

    Ok(RunningStack { join_set })
}

async fn monitor(join_set: &mut JoinSet<Result<()>>) -> Result<()> {
    while let Some(result) = join_set.join_next().await {
        match result {
            Ok(Ok(())) => {}
            Ok(Err(error)) => return Err(error),
            Err(join_error) => return Err(join_error).wrap_err("monitor join failed"),
        }
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn derive_namespace_prefixes_bare_robot_without_sanitizing() {
        assert_eq!(derive_namespace("42"), "/42");
        assert_eq!(derive_namespace("robot-01"), "/robot-01");
        assert_eq!(derive_namespace("robot//42"), "/robot//42");
        assert_eq!(derive_namespace("/robot/42"), "/robot/42");
        assert_eq!(derive_namespace("robot%01"), "/robot%01");
    }

    #[test]
    fn runtime_shutdown_timeout_does_not_wait_forever_for_blocking_tasks() {
        let (started_sender, started_receiver) = std::sync::mpsc::channel();
        let (release_sender, release_receiver) = std::sync::mpsc::channel::<()>();
        let started_at = std::time::Instant::now();

        let result = run_with_shutdown_timeout(
            async move {
                tokio::task::spawn_blocking(move || {
                    started_sender.send(()).expect("started signal should send");
                    let _ = release_receiver.recv();
                });
                started_receiver.recv().expect("blocking task should start");
            },
            std::time::Duration::from_millis(10),
        );

        drop(release_sender);
        result.expect("runtime should build and run");
        assert!(started_at.elapsed() < std::time::Duration::from_secs(1));
    }
}