mecha10-templates 0.8.2

//! Infrastructure template generation
//!
//! Handles generation of infrastructure files (docker-compose.yml, remote/rl Docker files,

use crate::paths;
use anyhow::{Context, Result};
use std::path::Path;

async fn read_project_template(assets_root: &Path, name: &str) -> Result<String> {
    let template_path = assets_root.join("project").join(name);
    tokio::fs::read_to_string(&template_path)
        .await
        .with_context(|| format!("Failed to read {:?} from downloaded mecha10-assets", template_path))
}

pub struct InfraTemplates;

impl InfraTemplates {
    /// Create a new InfraTemplates instance
    pub fn new() -> Self {
        Self
    }

    /// Create docker/ directory if it doesn't exist
    async fn ensure_docker_dir(&self, path: &Path) -> Result<()> {
        let docker_dir = path.join(paths::docker::DIR);
        if !docker_dir.exists() {
            tokio::fs::create_dir_all(&docker_dir).await?;
        }
        Ok(())
    }

    /// Create docker-compose.yml for the project, reading the template from the freshly-fetched
    /// `mecha10-assets` extraction at `assets_root`.
    pub async fn create_docker_compose(&self, path: &Path, project_name: &str, assets_root: &Path) -> Result<()> {
        self.ensure_docker_dir(path).await?;
        let template = read_project_template(assets_root, "docker-compose.yml.template").await?;
        let content = template.replace("{{project_name}}", project_name);
        tokio::fs::write(path.join(paths::docker::COMPOSE_FILE), content).await?;
        Ok(())
    }

    /// Create docker-compose.remote.yml, which runs the mecha10-remote container for AI/ML
    /// nodes with platform-specific binary dependencies, reading the template from the
    /// freshly-fetched `mecha10-assets` extraction at `assets_root`.
    pub async fn create_docker_compose_remote(&self, path: &Path, assets_root: &Path) -> Result<()> {
        self.ensure_docker_dir(path).await?;
        let framework_version = env!("CARGO_PKG_VERSION");
        let template = read_project_template(assets_root, "docker/docker-compose.remote.yml.template").await?;
        let content = template.replace("{{framework_version}}", framework_version);
        tokio::fs::write(path.join(paths::docker::COMPOSE_REMOTE_FILE), content).await?;
        Ok(())
    }

    /// Create Dockerfile.remote for building the mecha10-remote container. Users may customize
    /// it to add system dependencies for their AI nodes. Reads the template from the
    /// freshly-fetched `mecha10-assets` extraction at `assets_root`.
    pub async fn create_dockerfile_remote(&self, path: &Path, assets_root: &Path) -> Result<()> {
        self.ensure_docker_dir(path).await?;
        let framework_version = env!("CARGO_PKG_VERSION");
        let template = read_project_template(assets_root, "docker/Dockerfile.remote.template").await?;
        let content = template.replace("{{framework_version}}", framework_version);
        tokio::fs::write(path.join(paths::docker::DOCKERFILE_REMOTE), content).await?;
        Ok(())
    }

    /// Create docker/entrypoint.sh, reading the template from the freshly-fetched
    /// `mecha10-assets` extraction at `assets_root`.
    pub async fn create_entrypoint_remote(&self, path: &Path, assets_root: &Path) -> Result<()> {
        self.ensure_docker_dir(path).await?;
        let content = read_project_template(assets_root, "docker/entrypoint.remote.sh.template").await?;
        tokio::fs::write(path.join(paths::docker::ENTRYPOINT_REMOTE), content).await?;
        Ok(())
    }

    /// Create docker-compose.remote.yml, Dockerfile.remote, and entrypoint.sh, called during
    /// `mecha10 init` to set up remote node support.
    pub async fn create_remote_docker_files(&self, path: &Path, assets_root: &Path) -> Result<()> {
        self.ensure_docker_dir(path).await?;
        self.create_docker_compose_remote(path, assets_root).await?;
        self.create_dockerfile_remote(path, assets_root).await?;
        self.create_entrypoint_remote(path, assets_root).await?;
        Ok(())
    }

    /// Create docker-compose.rl.yml for the rl target - a trained RL policy in production.
    /// Modeled on `create_docker_compose_remote`, but excludes the simulator. Reads the template
    /// from the freshly-fetched `mecha10-assets` extraction at `assets_root`.
    pub async fn create_docker_compose_rl(&self, path: &Path, assets_root: &Path) -> Result<()> {
        self.ensure_docker_dir(path).await?;
        let framework_version = env!("CARGO_PKG_VERSION");
        let template = read_project_template(assets_root, "docker/docker-compose.rl.yml.template").await?;
        let content = template.replace("{{framework_version}}", framework_version);
        tokio::fs::write(path.join(paths::docker::COMPOSE_RL_FILE), content).await?;
        Ok(())
    }

    /// Create Dockerfile.rl for building the mecha10-rl container, reading the template from the
    /// freshly-fetched `mecha10-assets` extraction at `assets_root`.
    pub async fn create_dockerfile_rl(&self, path: &Path, assets_root: &Path) -> Result<()> {
        self.ensure_docker_dir(path).await?;
        let framework_version = env!("CARGO_PKG_VERSION");
        let template = read_project_template(assets_root, "docker/Dockerfile.rl.template").await?;
        let content = template.replace("{{framework_version}}", framework_version);
        tokio::fs::write(path.join(paths::docker::DOCKERFILE_RL), content).await?;
        Ok(())
    }

    /// Create docker/entrypoint-rl.sh, reading the template from the freshly-fetched
    /// `mecha10-assets` extraction at `assets_root`.
    pub async fn create_entrypoint_rl(&self, path: &Path, assets_root: &Path) -> Result<()> {
        self.ensure_docker_dir(path).await?;
        let content = read_project_template(assets_root, "docker/entrypoint-rl.sh.template").await?;
        tokio::fs::write(path.join(paths::docker::ENTRYPOINT_RL), content).await?;
        Ok(())
    }

    /// Create all rl target infrastructure files: docker-compose.rl.yml, Dockerfile.rl, and
    /// entrypoint-rl.sh. Called during `mecha10 init` for the rover template.
    pub async fn create_rl_docker_files(&self, path: &Path, assets_root: &Path) -> Result<()> {
        self.ensure_docker_dir(path).await?;
        self.create_docker_compose_rl(path, assets_root).await?;
        self.create_dockerfile_rl(path, assets_root).await?;
        self.create_entrypoint_rl(path, assets_root).await?;
        Ok(())
    }

    /// Create .env.example for the project, reading the template from the freshly-fetched
    /// `mecha10-assets` extraction at `assets_root`.
    pub async fn create_env_example(&self, path: &Path, project_name: &str, assets_root: &Path) -> Result<()> {
        let robot_id = project_name.replace('-', "_");

        let template = read_project_template(assets_root, ".env.example.template").await?;
        let content = template
            .replace("{{robot_id}}", &robot_id)
            .replace("{{project_name}}", project_name);

        tokio::fs::write(path.join(paths::env::EXAMPLE), &content).await?;

        tokio::fs::write(path.join(paths::env::FILE), &content).await?;

        Ok(())
    }
}

impl Default for InfraTemplates {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::{tempdir, TempDir};

    /// Stage a fake extracted `mecha10-assets` root containing the project-scaffolding
    /// templates these tests exercise, mirroring the real `mecha10-assets` layout
    /// (`project/...`, with Docker files nested at `project/docker/...`).
    fn stage_assets_root() -> TempDir {
        let assets_root = TempDir::new().unwrap();
        let docker_dir = assets_root.path().join("project/docker");
        std::fs::create_dir_all(&docker_dir).unwrap();

        std::fs::write(
            assets_root.path().join("project/docker-compose.yml.template"),
            "services:\n  app:\n    container_name: {{project_name}}\n",
        )
        .unwrap();
        std::fs::write(
            docker_dir.join("docker-compose.remote.yml.template"),
            "# framework {{framework_version}}\n",
        )
        .unwrap();
        std::fs::write(
            docker_dir.join("Dockerfile.remote.template"),
            format!(
                "# framework {{{{framework_version}}}}\nCOPY {} /entrypoint.sh\n",
                paths::docker::ENTRYPOINT_REMOTE
            ),
        )
        .unwrap();
        std::fs::write(
            docker_dir.join("entrypoint.remote.sh.template"),
            "#!/bin/sh\nexec mecha10-launcher start --target remote \"$@\"\n",
        )
        .unwrap();
        std::fs::write(
            docker_dir.join("docker-compose.rl.yml.template"),
            "# framework {{framework_version}}\n",
        )
        .unwrap();
        std::fs::write(
            docker_dir.join("Dockerfile.rl.template"),
            format!(
                "# framework {{{{framework_version}}}}\nCOPY {} /entrypoint.sh\n",
                paths::docker::ENTRYPOINT_RL
            ),
        )
        .unwrap();
        std::fs::write(
            docker_dir.join("entrypoint-rl.sh.template"),
            "#!/bin/sh\nexec mecha10-launcher start --target rl \"$@\"\n",
        )
        .unwrap();
        std::fs::write(
            assets_root.path().join("project/.env.example.template"),
            "ROBOT_ID={{robot_id}}\nMECHA10_PROJECT_NAME={{project_name}}\n",
        )
        .unwrap();

        assets_root
    }

    #[tokio::test]
    async fn create_remote_docker_files_writes_compose_dockerfile_and_entrypoint() {
        let dir = tempdir().unwrap();
        let assets_root = stage_assets_root();
        let infra = InfraTemplates::new();

        infra
            .create_remote_docker_files(dir.path(), assets_root.path())
            .await
            .unwrap();

        assert!(dir.path().join(paths::docker::COMPOSE_REMOTE_FILE).exists());
        assert!(dir.path().join(paths::docker::DOCKERFILE_REMOTE).exists());
        assert!(dir.path().join(paths::docker::ENTRYPOINT_REMOTE).exists());
    }

    #[tokio::test]
    async fn dockerfile_remote_copies_entrypoint_from_the_path_it_actually_writes_it_to() {
        let dir = tempdir().unwrap();
        let assets_root = stage_assets_root();
        let infra = InfraTemplates::new();

        infra
            .create_remote_docker_files(dir.path(), assets_root.path())
            .await
            .unwrap();

        let dockerfile = tokio::fs::read_to_string(dir.path().join(paths::docker::DOCKERFILE_REMOTE))
            .await
            .unwrap();
        assert!(
            dockerfile.contains(&format!("COPY {} /entrypoint.sh", paths::docker::ENTRYPOINT_REMOTE)),
            "Dockerfile.remote must COPY from {}, got:\n{dockerfile}",
            paths::docker::ENTRYPOINT_REMOTE
        );
    }

    #[tokio::test]
    async fn create_rl_docker_files_writes_compose_dockerfile_and_entrypoint() {
        let dir = tempdir().unwrap();
        let assets_root = stage_assets_root();
        let infra = InfraTemplates::new();

        infra
            .create_rl_docker_files(dir.path(), assets_root.path())
            .await
            .unwrap();

        assert!(dir.path().join(paths::docker::COMPOSE_RL_FILE).exists());
        assert!(dir.path().join(paths::docker::DOCKERFILE_RL).exists());
        assert!(dir.path().join(paths::docker::ENTRYPOINT_RL).exists());
    }

    #[tokio::test]
    async fn dockerfile_rl_copies_entrypoint_from_the_path_it_actually_writes_it_to() {
        let dir = tempdir().unwrap();
        let assets_root = stage_assets_root();
        let infra = InfraTemplates::new();

        infra
            .create_rl_docker_files(dir.path(), assets_root.path())
            .await
            .unwrap();

        let dockerfile = tokio::fs::read_to_string(dir.path().join(paths::docker::DOCKERFILE_RL))
            .await
            .unwrap();
        assert!(
            dockerfile.contains(&format!("COPY {} /entrypoint.sh", paths::docker::ENTRYPOINT_RL)),
            "Dockerfile.rl must COPY from {}, got:\n{dockerfile}",
            paths::docker::ENTRYPOINT_RL
        );
    }

    #[tokio::test]
    async fn rl_and_remote_docker_files_coexist_without_overwriting_each_other() {
        let dir = tempdir().unwrap();
        let assets_root = stage_assets_root();
        let infra = InfraTemplates::new();

        infra
            .create_remote_docker_files(dir.path(), assets_root.path())
            .await
            .unwrap();
        infra
            .create_rl_docker_files(dir.path(), assets_root.path())
            .await
            .unwrap();

        assert!(dir.path().join(paths::docker::ENTRYPOINT_REMOTE).exists());
        assert!(dir.path().join(paths::docker::ENTRYPOINT_RL).exists());
        assert_ne!(paths::docker::ENTRYPOINT_REMOTE, paths::docker::ENTRYPOINT_RL);

        let rl_entrypoint = tokio::fs::read_to_string(dir.path().join(paths::docker::ENTRYPOINT_RL))
            .await
            .unwrap();
        assert!(rl_entrypoint.contains("--target rl"));

        let remote_entrypoint = tokio::fs::read_to_string(dir.path().join(paths::docker::ENTRYPOINT_REMOTE))
            .await
            .unwrap();
        assert!(remote_entrypoint.contains("--target remote"));
    }

    #[tokio::test]
    async fn env_example_includes_project_name_distinct_from_robot_id() {
        let dir = tempdir().unwrap();
        let assets_root = stage_assets_root();
        let infra = InfraTemplates::new();

        infra
            .create_env_example(dir.path(), "my-robot", assets_root.path())
            .await
            .unwrap();

        let content = tokio::fs::read_to_string(dir.path().join(paths::env::EXAMPLE))
            .await
            .unwrap();
        assert!(content.contains("ROBOT_ID=my_robot"));
        assert!(content.contains("MECHA10_PROJECT_NAME=my-robot"));
    }
}