mecha10-cli 0.9.0

//! RL project bundle download (LAB-2769)
//!
//! Downloads and unpacks a robot-scoped RL project bundle - a tar.gz of a project's
//! `configs/rl/` and `configs/simulation/{models,environments}` - before training starts.

use anyhow::{Context, Result};
use flate2::read::GzDecoder;
use std::path::Path;
use tar::Archive;

/// URL for LAB-2767's robot-scoped RL project bundle upload endpoint.
///
/// Assumed shape: `PUT {control_plane_url}/api/robots/{robot_id}/rl-bundle`.
fn build_rl_bundle_url(control_plane_url: &str, robot_id: &str) -> String {
    format!(
        "{}/api/robots/{}/rl-bundle",
        control_plane_url.trim_end_matches('/'),
        robot_id
    )
}

/// Download the RL project bundle previously uploaded for `robot_id`, returning the raw tar.gz
/// bytes for [`unpack_rl_bundle_archive`] to unpack.
pub async fn download_rl_bundle(
    control_plane_url: &str,
    robot_id: &str,
    api_key: &str,
    user_id: &str,
) -> Result<Vec<u8>> {
    let url = build_rl_bundle_url(control_plane_url, robot_id);

    let client = reqwest::Client::new();
    let response = client
        .get(&url)
        .header("Authorization", format!("Bearer {}", api_key))
        .header("X-User-ID", user_id)
        .send()
        .await
        .context("Failed to reach mecha10-api")?;

    let status = response.status();
    if !status.is_success() {
        let body = response.text().await.unwrap_or_default();
        anyhow::bail!("RL bundle download failed with status {}: {}", status, body);
    }

    let bytes = response
        .bytes()
        .await
        .context("Failed to read RL bundle download response body")?;

    Ok(bytes.to_vec())
}

/// Unpack an RL project bundle archive (as returned by [`download_rl_bundle`]) into `dest_root`.
/// Overwrites files already present at the destination; does not delete anything at the
/// destination that's absent from the archive.
pub fn unpack_rl_bundle_archive(archive_bytes: &[u8], dest_root: &Path) -> Result<()> {
    std::fs::create_dir_all(dest_root)
        .with_context(|| format!("Failed to create destination directory {}", dest_root.display()))?;

    let decoder = GzDecoder::new(archive_bytes);
    let mut archive = Archive::new(decoder);
    archive
        .unpack(dest_root)
        .context("Failed to unpack RL project bundle")?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use flate2::write::GzEncoder;
    use flate2::Compression;
    use tempfile::tempdir;

    /// Builds a minimal tar.gz archive containing a single `configs/rl/jobs/train.yml` entry.
    fn sample_bundle_archive() -> Vec<u8> {
        let mut bytes = Vec::new();
        let encoder = GzEncoder::new(&mut bytes, Compression::default());
        let mut builder = tar::Builder::new(encoder);

        let src = tempdir().unwrap();
        std::fs::create_dir_all(src.path().join("configs/rl/jobs")).unwrap();
        std::fs::write(src.path().join("configs/rl/jobs/train.yml"), "flow_type: train").unwrap();
        builder
            .append_dir_all("configs/rl", src.path().join("configs/rl"))
            .unwrap();

        let encoder = builder.into_inner().unwrap();
        encoder.finish().unwrap();
        bytes
    }

    #[test]
    fn build_rl_bundle_url_scopes_by_robot_id_and_strips_trailing_slash() {
        assert_eq!(
            build_rl_bundle_url("https://api.example.com/", "robot-1"),
            "https://api.example.com/api/robots/robot-1/rl-bundle"
        );
    }

    #[tokio::test]
    async fn download_rl_bundle_errors_when_server_unreachable() {
        let result = download_rl_bundle("http://127.0.0.1:9", "robot-1", "test-api-key", "user-1").await;
        assert!(result.is_err());
    }

    #[test]
    fn unpack_rl_bundle_archive_unpacks_entries_into_the_destination() {
        let bytes = sample_bundle_archive();

        let dest = tempdir().unwrap();
        unpack_rl_bundle_archive(&bytes, dest.path()).unwrap();

        assert_eq!(
            std::fs::read_to_string(dest.path().join("configs/rl/jobs/train.yml")).unwrap(),
            "flow_type: train"
        );
    }

    #[test]
    fn unpack_rl_bundle_archive_creates_missing_destination_directory() {
        let bytes = sample_bundle_archive();

        let dest_parent = tempdir().unwrap();
        let dest = dest_parent.path().join("does-not-exist-yet");
        assert!(!dest.exists());

        unpack_rl_bundle_archive(&bytes, &dest).unwrap();

        assert!(dest.join("configs/rl/jobs/train.yml").exists());
    }

    #[test]
    fn unpack_rl_bundle_archive_errors_on_invalid_archive_bytes() {
        let dest = tempdir().unwrap();
        let result = unpack_rl_bundle_archive(b"not a gzip archive", dest.path());
        assert!(result.is_err());
    }
}