mecha10-cli 0.8.0

//! Job catalog for `mecha10 train` (LAB-1776).
//!
//! A "job" is a YAML file under a project's `configs/rl/jobs/` directory that maps to one of
//! the Kestra `flows/ml-pipeline/*.yml` flow types plus an env/algorithm/hyperparameter selection.

use anyhow::{bail, Context, Result};
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};

pub const JOB_CATALOG_DIR: &str = "configs/rl/jobs";

pub use mecha10_job_validation::FlowType;

/// One dimension of a hyperparameter search space for a tuning job.
#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, schemars::JsonSchema)]
pub struct SearchSpaceEntry {
    pub name: String,
    #[serde(rename = "type")]
    pub type_: SearchSpaceType,
    #[serde(default)]
    pub low: Option<f64>,
    #[serde(default)]
    pub high: Option<f64>,
    #[serde(default)]
    pub log: Option<bool>,
    #[serde(default)]
    pub choices: Option<Vec<serde_json::Value>>,
}

/// `search_space` entry kind - `float`/`int` take `low`/`high` (and optional `log`),
/// `categorical` takes `choices`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Serialize, schemars::JsonSchema)]
#[serde(rename_all = "lowercase")]
pub enum SearchSpaceType {
    Float,
    Int,
    Categorical,
}

/// A training job definition, loaded from `configs/rl/jobs/<name>.yml`.
#[derive(Debug, Clone, Deserialize, schemars::JsonSchema)]
#[allow(dead_code)]
pub struct JobDefinition {
    pub name: String,
    #[serde(default)]
    pub description: Option<String>,
    pub flow_type: FlowType,
    /// Name of the simulation environment to train in. If omitted, it's resolved from the
    /// selected task.
    #[serde(default)]
    #[schemars(example = Some("warehouse_nav".to_string()))]
    pub env_name: Option<String>,
    /// Name of the task/curriculum to train on, selected from the project's task catalog.
    #[serde(default)]
    #[schemars(example = Some("patrol_curriculum".to_string()))]
    pub task_name: Option<String>,
    /// Name of the reinforcement learning algorithm to use for this job.
    #[serde(default)]
    #[schemars(example = Some("ppo".to_string()))]
    pub algorithm: Option<String>,

    /// Name of the policy network architecture to use. If omitted, it's auto-detected from the
    /// environment.
    #[serde(default)]
    #[schemars(example = Some("MlpPolicy".to_string()))]
    #[allow(dead_code, reason = "schema/deserialize-only, see doc comment above")]
    pub policy: Option<String>,
    #[serde(default)]
    #[allow(dead_code, reason = "schema/deserialize-only, see `policy`'s doc comment above")]
    pub seed: Option<i64>,
    #[serde(default)]
    #[allow(dead_code, reason = "schema/deserialize-only, see `policy`'s doc comment above")]
    pub environment_map: Option<String>,
    #[serde(default)]
    pub total_timesteps: Option<i64>,
    #[serde(default)]
    pub learning_rate: Option<f64>,
    #[serde(default)]
    #[allow(dead_code, reason = "schema/deserialize-only, see `policy`'s doc comment above")]
    pub linear_lr_schedule: Option<bool>,
    #[serde(default)]
    pub n_steps: Option<i64>,
    #[serde(default)]
    #[allow(dead_code, reason = "schema/deserialize-only, see `policy`'s doc comment above")]
    pub batch_size: Option<i64>,
    #[serde(default)]
    pub n_epochs: Option<i64>,
    #[serde(default)]
    #[allow(dead_code, reason = "schema/deserialize-only, see `policy`'s doc comment above")]
    pub gamma: Option<f64>,
    #[serde(default)]
    #[allow(dead_code, reason = "schema/deserialize-only, see `policy`'s doc comment above")]
    pub gae_lambda: Option<f64>,
    #[serde(default)]
    #[allow(dead_code, reason = "schema/deserialize-only, see `policy`'s doc comment above")]
    pub clip_range: Option<f64>,
    #[serde(default)]
    #[allow(dead_code, reason = "schema/deserialize-only, see `policy`'s doc comment above")]
    pub ent_coef: Option<f64>,
    #[serde(default)]
    #[allow(dead_code, reason = "schema/deserialize-only, see `policy`'s doc comment above")]
    pub vf_coef: Option<f64>,
    #[serde(default)]
    #[allow(dead_code, reason = "schema/deserialize-only, see `policy`'s doc comment above")]
    pub max_grad_norm: Option<f64>,
    /// Number of evaluations with no improvement to wait before stopping training early.
    #[serde(default)]
    #[schemars(example = Some(10))]
    pub patience: Option<i64>,
    #[serde(default)]
    pub min_delta: Option<f64>,
    /// Number of parallel environments to run during training.
    #[serde(default)]
    #[schemars(example = Some(8))]
    pub n_envs: Option<i64>,

    #[serde(default)]
    pub n_trials: Option<i64>,
    #[serde(default)]
    pub tune_timesteps: Option<i64>,
    #[serde(default)]
    pub eval_eps: Option<i64>,
    /// Hyperparameter search space to explore for a tuning job. If omitted, a default search
    /// space is used.
    #[serde(default)]
    pub search_space: Option<Vec<SearchSpaceEntry>>,

    #[serde(skip)]
    #[schemars(skip)]
    pub source_path: PathBuf,
}

#[derive(Debug, Clone, Serialize)]
pub struct CreateJobRequest {
    pub name: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,
    pub algorithm: String,
    pub robot_id: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub project_id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub project_name: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub env_name: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub task_name: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub total_timesteps: Option<i64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub n_steps: Option<i64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub n_epochs: Option<i64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub learning_rate: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub patience: Option<i64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub min_delta: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub n_envs: Option<i64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub config_yaml: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub resume_run_id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub n_trials: Option<i64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tune_timesteps: Option<i64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub eval_eps: Option<i64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub search_space: Option<Vec<SearchSpaceEntry>>,
}

/// Read `configs/rl/tasks/<task_name>/task.json`'s `env_name` field, relative to `project_root`.
/// Returns `None` on any failure (missing file, invalid JSON, missing/non-string field).
fn env_name_from_task_json(project_root: &Path, task_name: &str) -> Option<String> {
    task_json_field(project_root, task_name, "env_name")
}

fn task_json_field(project_root: &Path, task_name: &str, field: &str) -> Option<String> {
    let path = project_root
        .join(mecha10_cli_core::paths::project::RL_TASKS_DIR)
        .join(task_name)
        .join("task.json");
    let content = std::fs::read_to_string(path).ok()?;
    let value: serde_json::Value = serde_json::from_str(&content).ok()?;
    value.get(field)?.as_str().map(str::to_string)
}

/// Drop `value` if it equals `sentinel` - the "unset" placeholder some YAML/env fields use
/// instead of a real `null` (e.g. `-1` / `-1.0`).
fn non_sentinel<T: PartialEq>(value: Option<T>, sentinel: T) -> Option<T> {
    value.filter(|v| *v != sentinel)
}

fn non_sentinel_i64(value: Option<i64>) -> Option<i64> {
    non_sentinel(value, -1)
}

fn non_sentinel_f64(value: Option<f64>) -> Option<f64> {
    non_sentinel(value, -1.0)
}

impl JobDefinition {
    /// This job's fields in `mecha10-job-validation`'s shared vocabulary (see
    /// [`mecha10_job_validation::JobFields`]).
    fn validation_fields(&self, resume_run_id: Option<&str>) -> mecha10_job_validation::JobFields {
        mecha10_job_validation::JobFields {
            task_name: self.task_name.clone(),
            algorithm: self.algorithm.clone(),
            total_timesteps: non_sentinel_i64(self.total_timesteps),
            resume_run_id: resume_run_id.map(str::to_string),
            tune_run_id: None,
            n_trials: non_sentinel_i64(self.n_trials),
            tune_timesteps: non_sentinel_i64(self.tune_timesteps),
            eval_eps: non_sentinel_i64(self.eval_eps),
        }
    }

    /// Validate this job against `mecha10-job-validation`'s flow_type requirement table, plus a
    /// `task_name`-exists check against the project's own task catalog.
    pub fn validate(&self, project_root: &Path, resume_run_id: Option<&str>) -> Result<()> {
        let missing = mecha10_job_validation::validate_job(self.flow_type, &self.validation_fields(resume_run_id));
        if !missing.is_empty() {
            let details = missing.iter().map(|m| m.to_string()).collect::<Vec<_>>().join("; ");
            bail!("job '{}' is missing required fields: {details}", self.name);
        }

        let mut missing_hyperparams = Vec::new();
        if non_sentinel_i64(self.total_timesteps).is_none() {
            missing_hyperparams.push("total_timesteps");
        }
        if self.flow_type != FlowType::ResumeRun && non_sentinel_i64(self.n_steps).is_none() {
            missing_hyperparams.push("n_steps");
        }
        if !missing_hyperparams.is_empty() {
            bail!(
                "job '{}' is missing required field(s): {}",
                self.name,
                missing_hyperparams.join(", ")
            );
        }

        if let Some(task_name) = self.task_name.as_deref() {
            mecha10_job_validation::task_exists(project_root, task_name)
                .with_context(|| format!("job '{}' has an invalid task_name", self.name))?;
        }

        Ok(())
    }

    /// Resolve which `mecha10_rl.envs` registry name this job targets: an explicit `env_name`
    /// wins, else it's read from `task_name`'s own `task.json`. `None` when neither is available.
    pub fn resolve_env_name(&self, project_root: &Path) -> Option<String> {
        self.env_name.clone().or_else(|| {
            self.task_name
                .as_deref()
                .and_then(|task_name| env_name_from_task_json(project_root, task_name))
        })
    }

    /// Whether this job can be dispatched remotely with today's `CreateJobRequest`/`create_job`
    /// contract. `tune-then-train`/`train-from-tune` aren't dispatchable yet and stay local-only.
    pub fn supports_remote(&self) -> bool {
        matches!(self.flow_type, FlowType::Train | FlowType::ResumeRun | FlowType::Tune)
    }

    /// Build the remote dispatch request body. Only meaningful when `supports_remote()` is true;
    /// callers should check that first.
    pub fn to_create_job_request(
        &self,
        project_root: &Path,
        robot_id: &str,
        project_name: &str,
        resume_run_id: Option<&str>,
    ) -> Result<CreateJobRequest> {
        let config_yaml = std::fs::read_to_string(&self.source_path).with_context(|| {
            format!(
                "failed to read {} - this job definition's own file must exist before dispatching remotely",
                self.source_path.display()
            )
        })?;

        Ok(CreateJobRequest {
            name: self.name.clone(),
            description: self.description.clone(),
            algorithm: self.algorithm.clone().unwrap_or_default(),
            robot_id: robot_id.to_string(),
            project_id: Some(robot_id.to_string()),
            project_name: Some(project_name.to_string()),
            env_name: self.resolve_env_name(project_root),
            task_name: self.task_name.clone(),
            total_timesteps: non_sentinel_i64(self.total_timesteps),
            n_steps: non_sentinel_i64(self.n_steps),
            n_epochs: non_sentinel_i64(self.n_epochs),
            learning_rate: non_sentinel_f64(self.learning_rate),
            patience: non_sentinel_i64(self.patience),
            min_delta: non_sentinel_f64(self.min_delta),
            n_envs: non_sentinel_i64(self.n_envs),
            config_yaml: Some(config_yaml),
            resume_run_id: resume_run_id
                .map(str::trim)
                .filter(|id| !id.is_empty())
                .map(str::to_string),
            n_trials: non_sentinel_i64(self.n_trials),
            tune_timesteps: non_sentinel_i64(self.tune_timesteps),
            eval_eps: non_sentinel_i64(self.eval_eps),
            search_space: self.search_space.clone(),
        })
    }
}

/// Load every `.yml`/`.yaml` job definition under `{project_root}/configs/rl/jobs/`, sorted by
/// file name for a stable menu order. Returns an empty list (not an error) when the directory
/// doesn't exist - `mecha10 train` treats that as "no jobs configured yet" and points the user at
/// `mecha10 init`/the bundled defaults rather than failing.
pub fn load_job_catalog(project_root: &Path) -> Result<Vec<JobDefinition>> {
    let jobs_dir = project_root.join(JOB_CATALOG_DIR);
    if !jobs_dir.exists() {
        return Ok(Vec::new());
    }

    let mut paths: Vec<PathBuf> = std::fs::read_dir(&jobs_dir)
        .with_context(|| format!("failed to read {}", jobs_dir.display()))?
        .filter_map(|entry| entry.ok())
        .map(|entry| entry.path())
        .filter(|path| {
            path.extension()
                .and_then(|ext| ext.to_str())
                .is_some_and(|ext| ext == "yml" || ext == "yaml")
        })
        .collect();
    paths.sort();

    let mut jobs = Vec::with_capacity(paths.len());
    for path in paths {
        let content = std::fs::read_to_string(&path).with_context(|| format!("failed to read {}", path.display()))?;
        let mut job: JobDefinition = serde_yaml::from_str(&content)
            .with_context(|| format!("failed to parse job definition {}", path.display()))?;
        job.source_path = path;
        jobs.push(job);
    }

    Ok(jobs)
}

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

    fn write_job(dir: &Path, file_name: &str, contents: &str) {
        std::fs::write(dir.join(file_name), contents).unwrap();
    }

    fn write_task_json(project_root: &Path, task_name: &str, contents: &str) {
        let task_dir = project_root
            .join(mecha10_cli_core::paths::project::RL_TASKS_DIR)
            .join(task_name);
        std::fs::create_dir_all(&task_dir).unwrap();
        std::fs::write(task_dir.join("task.json"), contents).unwrap();
    }

    #[test]
    fn resolve_env_name_prefers_an_explicit_env_name_over_the_tasks_own() {
        let dir = TempDir::new().unwrap();
        write_task_json(dir.path(), "find_target", r#"{"env_name": "basic_arena"}"#);

        let job = JobDefinition {
            env_name: Some("rover_arena".to_string()),
            task_name: Some("find_target".to_string()),
            ..base_job()
        };
        assert_eq!(job.resolve_env_name(dir.path()), Some("rover_arena".to_string()));
    }

    #[test]
    fn resolve_env_name_falls_back_to_the_tasks_own_task_json() {
        let dir = TempDir::new().unwrap();
        write_task_json(
            dir.path(),
            "find_target",
            r#"{"env_name": "basic_arena", "model": "humanoid"}"#,
        );

        let job = JobDefinition {
            task_name: Some("find_target".to_string()),
            ..base_job()
        };
        assert_eq!(job.resolve_env_name(dir.path()), Some("basic_arena".to_string()));
    }

    #[test]
    fn resolve_env_name_is_none_when_task_has_no_env_name_field() {
        let dir = TempDir::new().unwrap();
        write_task_json(dir.path(), "find_target", r#"{"model": "humanoid"}"#);

        let job = JobDefinition {
            task_name: Some("find_target".to_string()),
            ..base_job()
        };
        assert_eq!(job.resolve_env_name(dir.path()), None);
    }

    #[test]
    fn resolve_env_name_is_none_without_an_env_name_or_a_task_name() {
        let dir = TempDir::new().unwrap();
        let job = base_job();
        assert_eq!(job.resolve_env_name(dir.path()), None);
    }

    /// A minimal, fully-`None`/default `JobDefinition` for tests to `..base_job()` from, so each
    /// test only spells out the fields it actually cares about.
    fn base_job() -> JobDefinition {
        JobDefinition {
            name: "job".to_string(),
            description: None,
            flow_type: FlowType::Train,
            env_name: None,
            task_name: None,
            algorithm: Some("ppo".to_string()),
            policy: None,
            seed: None,
            environment_map: None,
            total_timesteps: None,
            learning_rate: None,
            linear_lr_schedule: None,
            n_steps: None,
            batch_size: None,
            n_epochs: None,
            gamma: None,
            gae_lambda: None,
            clip_range: None,
            ent_coef: None,
            vf_coef: None,
            max_grad_norm: None,
            patience: None,
            min_delta: None,
            n_envs: None,
            n_trials: None,
            tune_timesteps: None,
            eval_eps: None,
            search_space: None,
            source_path: PathBuf::new(),
        }
    }

    #[test]
    fn load_job_catalog_returns_empty_when_directory_missing() {
        let temp = TempDir::new().unwrap();
        let jobs = load_job_catalog(temp.path()).unwrap();
        assert!(jobs.is_empty());
    }

    #[test]
    fn load_job_catalog_parses_yaml_files_sorted_by_name() {
        let temp = TempDir::new().unwrap();
        let jobs_dir = temp.path().join(JOB_CATALOG_DIR);
        std::fs::create_dir_all(&jobs_dir).unwrap();

        write_job(
            &jobs_dir,
            "train.yml",
            "name: train\nflow_type: train\nalgorithm: ppo\ntotal_timesteps: 1000000\n",
        );
        write_job(
            &jobs_dir,
            "hypertune.yml",
            "name: hypertune\nflow_type: tune\nalgorithm: ppo\nn_trials: 10\n",
        );

        let jobs = load_job_catalog(temp.path()).unwrap();

        assert_eq!(jobs.len(), 2);
        assert_eq!(jobs[0].name, "hypertune");
        assert_eq!(jobs[0].flow_type, FlowType::Tune);
        assert_eq!(jobs[1].name, "train");
        assert_eq!(jobs[1].flow_type, FlowType::Train);
        assert_eq!(jobs[1].total_timesteps, Some(1000000));
    }

    #[test]
    fn load_job_catalog_ignores_non_yaml_files() {
        let temp = TempDir::new().unwrap();
        let jobs_dir = temp.path().join(JOB_CATALOG_DIR);
        std::fs::create_dir_all(&jobs_dir).unwrap();

        write_job(&jobs_dir, "README.md", "not a job");
        write_job(
            &jobs_dir,
            "train.yml",
            "name: train\nflow_type: train\nalgorithm: ppo\n",
        );

        let jobs = load_job_catalog(temp.path()).unwrap();
        assert_eq!(jobs.len(), 1);
    }

    #[test]
    fn supports_remote_for_train_resume_run_and_tune() {
        let base = base_job();
        assert!(base.supports_remote());

        let resume = JobDefinition {
            flow_type: FlowType::ResumeRun,
            ..base.clone()
        };
        assert!(resume.supports_remote());

        let tune = JobDefinition {
            flow_type: FlowType::Tune,
            ..base.clone()
        };
        assert!(tune.supports_remote());

        let tune_then_train = JobDefinition {
            flow_type: FlowType::TuneThenTrain,
            ..base.clone()
        };
        assert!(!tune_then_train.supports_remote());

        let train_from_tune = JobDefinition {
            flow_type: FlowType::TrainFromTune,
            ..base
        };
        assert!(!train_from_tune.supports_remote());
    }

    #[test]
    fn to_create_job_request_forwards_resume_run_id() {
        let temp = TempDir::new().unwrap();
        let jobs_dir = temp.path().join("configs/rl/jobs");
        std::fs::create_dir_all(&jobs_dir).unwrap();
        let source_path = jobs_dir.join("train-from-checkpoint.yml");
        std::fs::write(&source_path, "name: train-from-checkpoint\nflow_type: resume-run\n").unwrap();

        let job = JobDefinition {
            name: "train-from-checkpoint".to_string(),
            flow_type: FlowType::ResumeRun,
            source_path,
            ..base_job()
        };

        let req = job
            .to_create_job_request(temp.path(), "my-project", "my-robot", Some("  run-abc123  "))
            .unwrap();

        assert_eq!(req.resume_run_id.as_deref(), Some("run-abc123"));
    }

    #[test]
    fn to_create_job_request_omits_empty_resume_run_id() {
        let temp = TempDir::new().unwrap();
        let jobs_dir = temp.path().join("configs/rl/jobs");
        std::fs::create_dir_all(&jobs_dir).unwrap();
        let source_path = jobs_dir.join("train-from-checkpoint.yml");
        std::fs::write(&source_path, "name: train-from-checkpoint\nflow_type: resume-run\n").unwrap();

        let job = JobDefinition {
            name: "train-from-checkpoint".to_string(),
            flow_type: FlowType::ResumeRun,
            source_path,
            ..base_job()
        };

        let req = job
            .to_create_job_request(temp.path(), "my-project", "my-robot", Some("   "))
            .unwrap();

        assert_eq!(req.resume_run_id, None);
    }

    /// Contract test against `mecha10-job-validation`'s `JobField` list: every field applicable
    /// to a remotely-dispatchable flow type must have a matching wire key on this struct.
    #[test]
    fn create_job_request_has_a_field_for_every_dispatchable_job_field() {
        let dispatchable = [FlowType::Train, FlowType::ResumeRun, FlowType::Tune];

        let req = CreateJobRequest {
            name: "job".to_string(),
            description: None,
            algorithm: "ppo".to_string(),
            robot_id: "robot-1".to_string(),
            project_id: None,
            project_name: None,
            env_name: None,
            task_name: Some("find_target".to_string()),
            total_timesteps: Some(1_000_000),
            n_steps: None,
            n_epochs: None,
            learning_rate: None,
            patience: None,
            min_delta: None,
            n_envs: None,
            config_yaml: None,
            resume_run_id: Some("run-123".to_string()),
            n_trials: Some(10),
            tune_timesteps: Some(500_000),
            eval_eps: Some(5),
            search_space: None,
        };
        let json = serde_json::to_value(&req).expect("CreateJobRequest serializes");

        for field in mecha10_job_validation::all_fields() {
            let applicable = dispatchable
                .iter()
                .any(|&ft| field.requirement(ft) != mecha10_job_validation::FieldRequirement::NotApplicable);
            if applicable {
                assert!(
                    json.get(field.wire_name()).is_some(),
                    "CreateJobRequest has no '{}' field for {field:?}, which is applicable to a \
                     dispatchable flow type",
                    field.wire_name()
                );
            }
        }
    }

    #[test]
    fn to_create_job_request_forwards_tune_fields() {
        let temp = TempDir::new().unwrap();
        let jobs_dir = temp.path().join("configs/rl/jobs");
        std::fs::create_dir_all(&jobs_dir).unwrap();
        let source_path = jobs_dir.join("hypertune.yml");
        std::fs::write(&source_path, "name: hypertune\nflow_type: tune\n").unwrap();

        let job = JobDefinition {
            name: "hypertune".to_string(),
            flow_type: FlowType::Tune,
            n_trials: Some(20),
            tune_timesteps: Some(250_000),
            eval_eps: Some(3),
            source_path,
            ..base_job()
        };

        let req = job
            .to_create_job_request(temp.path(), "my-project", "my-robot", None)
            .unwrap();

        assert_eq!(req.n_trials, Some(20));
        assert_eq!(req.tune_timesteps, Some(250_000));
        assert_eq!(req.eval_eps, Some(3));
    }

    #[test]
    fn to_create_job_request_forwards_n_envs() {
        let temp = TempDir::new().unwrap();
        let jobs_dir = temp.path().join("configs/rl/jobs");
        std::fs::create_dir_all(&jobs_dir).unwrap();
        let source_path = jobs_dir.join("train.yml");
        std::fs::write(&source_path, "name: train\nflow_type: train\nn_envs: 8\n").unwrap();

        let job = JobDefinition {
            name: "train".to_string(),
            n_envs: Some(8),
            source_path,
            ..base_job()
        };

        let req = job
            .to_create_job_request(temp.path(), "my-project", "my-robot", None)
            .unwrap();

        assert_eq!(req.n_envs, Some(8));
    }

    #[test]
    fn to_create_job_request_forwards_search_space() {
        let temp = TempDir::new().unwrap();
        let jobs_dir = temp.path().join("configs/rl/jobs");
        std::fs::create_dir_all(&jobs_dir).unwrap();
        let source_path = jobs_dir.join("hypertune.yml");
        let yaml = "name: hypertune\nflow_type: tune\nsearch_space:\n  - name: learning_rate\n    \
                    type: float\n    low: 1e-5\n    high: 1e-3\n    log: true\n  - name: n_steps\n    \
                    type: categorical\n    choices: [1024, 2048, 4096]\n  - name: n_epochs\n    \
                    type: int\n    low: 5\n    high: 20\n";
        std::fs::write(&source_path, yaml).unwrap();
        let job: JobDefinition = serde_yaml::from_str(yaml).unwrap();
        let job = JobDefinition { source_path, ..job };

        let req = job
            .to_create_job_request(temp.path(), "my-project", "my-robot", None)
            .unwrap();

        let search_space = req.search_space.expect("search_space forwarded");
        assert_eq!(search_space.len(), 3);
        assert_eq!(search_space[0].name, "learning_rate");
        assert_eq!(search_space[0].type_, SearchSpaceType::Float);
        assert_eq!(search_space[0].low, Some(1e-5));
        assert_eq!(search_space[0].high, Some(1e-3));
        assert_eq!(search_space[0].log, Some(true));
        assert_eq!(search_space[1].name, "n_steps");
        assert_eq!(search_space[1].type_, SearchSpaceType::Categorical);
        assert_eq!(
            search_space[1].choices,
            Some(vec![
                serde_json::json!(1024),
                serde_json::json!(2048),
                serde_json::json!(4096)
            ])
        );
        assert_eq!(search_space[2].name, "n_epochs");
        assert_eq!(search_space[2].type_, SearchSpaceType::Int);
        assert_eq!(search_space[2].low, Some(5.0));
        assert_eq!(search_space[2].high, Some(20.0));
    }

    #[test]
    fn to_create_job_request_omits_search_space_when_unset() {
        let temp = TempDir::new().unwrap();
        let jobs_dir = temp.path().join("configs/rl/jobs");
        std::fs::create_dir_all(&jobs_dir).unwrap();
        let source_path = jobs_dir.join("train.yml");
        std::fs::write(&source_path, "name: train\nflow_type: train\n").unwrap();

        let job = JobDefinition {
            name: "train".to_string(),
            source_path,
            ..base_job()
        };

        let req = job
            .to_create_job_request(temp.path(), "my-project", "my-robot", None)
            .unwrap();

        assert_eq!(req.search_space, None);
    }

    #[test]
    fn to_create_job_request_drops_sentinel_values() {
        let temp = TempDir::new().unwrap();
        let jobs_dir = temp.path().join("configs/rl/jobs");
        std::fs::create_dir_all(&jobs_dir).unwrap();
        let source_path = jobs_dir.join("train.yml");
        std::fs::write(&source_path, "name: train\nflow_type: train\nenv_name: rover\n").unwrap();

        let job = JobDefinition {
            name: "train".to_string(),
            description: Some("desc".to_string()),
            env_name: Some("rover".to_string()),
            total_timesteps: Some(-1),
            n_steps: Some(2048),
            n_epochs: Some(-1),
            learning_rate: Some(-1.0),
            patience: Some(50),
            min_delta: Some(-1.0),
            n_envs: Some(-1),
            source_path,
            ..base_job()
        };

        let req = job
            .to_create_job_request(temp.path(), "my-project", "my-robot", None)
            .unwrap();

        assert_eq!(req.name, "train");
        assert_eq!(req.algorithm, "ppo");
        assert_eq!(req.robot_id, "my-project");
        assert_eq!(req.project_id.as_deref(), Some("my-project"));
        assert_eq!(req.project_name.as_deref(), Some("my-robot"));
        assert_eq!(req.total_timesteps, None);
        assert_eq!(req.n_steps, Some(2048));
        assert_eq!(req.n_epochs, None);
        assert_eq!(req.learning_rate, None);
        assert_eq!(req.patience, Some(50));
        assert_eq!(req.min_delta, None);
        assert_eq!(req.n_envs, None);
        assert_eq!(
            req.config_yaml.as_deref(),
            Some("name: train\nflow_type: train\nenv_name: rover\n")
        );
        assert_eq!(req.task_name, None);
    }

    #[test]
    fn to_create_job_request_forwards_task_name() {
        let temp = TempDir::new().unwrap();
        let jobs_dir = temp.path().join("configs/rl/jobs");
        std::fs::create_dir_all(&jobs_dir).unwrap();
        let source_path = jobs_dir.join("train.yml");
        std::fs::write(&source_path, "name: train\nflow_type: train\n").unwrap();

        let job = JobDefinition {
            name: "train".to_string(),
            task_name: Some("navigate".to_string()),
            source_path,
            ..base_job()
        };

        let req = job
            .to_create_job_request(temp.path(), "my-project", "my-robot", None)
            .unwrap();

        assert_eq!(req.task_name.as_deref(), Some("navigate"));
    }

    #[test]
    fn to_create_job_request_errors_when_the_jobs_own_file_is_missing() {
        let temp = TempDir::new().unwrap();
        let job = JobDefinition {
            name: "train".to_string(),
            source_path: temp.path().join("configs/rl/jobs/train.yml"),
            ..base_job()
        };

        let err = job
            .to_create_job_request(temp.path(), "my-project", "my-robot", None)
            .unwrap_err();
        assert!(err.to_string().contains("configs/rl/jobs/train.yml"));
    }
}