mecha10-cli 0.8.0

//! `mecha10 train` command handler.
//!
//! Interactive flow: validate configs, pick a job and task, push configs, then dispatch the
//! job remotely via mecha10-api's job-creation endpoint.

use crate::commands::TrainArgs;
use crate::config_validation::{validate_configs, ConfigKind};
use crate::context::CliContext;
use crate::handlers::config::push_configs;
use crate::services::credentials::CredentialsService;
use crate::services::train_jobs::{load_job_catalog, FlowType, JobDefinition};
use crate::ui::style;
use anyhow::{bail, Context, Result};
use dialoguer::{theme::ColorfulTheme, Select};
use serde::Deserialize;
use std::path::Path;

const JOB_LIST_PATH: &str = "/api/model-deployments/jobs";

const JOB_CREATE_PATH: &str = "/api/model-deployments/jobs";

/// Handle `mecha10 train`.
pub async fn handle_train(ctx: &mut CliContext, args: &TrainArgs) -> Result<()> {
    println!();
    println!("{}", style::heading("Mecha10 Training"));
    println!("{}", style::separator(47));
    println!();

    if !ctx.is_project_initialized() {
        bail!("Not inside a mecha10 project (no mecha10.json found). Run `mecha10 init` first.");
    }

    let assets_dir = ctx
        .assets_download()
        .fetch_and_extract()
        .await
        .context("Failed to download mecha10-assets for config validation")?;
    let schemas_dir = assets_dir.path().join("configs/simulation/schemas");
    validate_configs(
        &ctx.working_dir,
        &[
            ConfigKind::RlJobs,
            ConfigKind::SimulationEnvironments,
            ConfigKind::ProjectAndNodes,
        ],
        Some(&schemas_dir),
    )
    .context("pre-flight config validation failed for `mecha10 train`")?;

    let jobs = load_job_catalog(&ctx.working_dir).context("failed to load job catalog from configs/rl/jobs/")?;

    if jobs.is_empty() {
        println!("No training jobs found under configs/rl/jobs/.");
        println!();
        println!("A freshly scaffolded rover project ships a standard set (train,");
        println!("train-from-checkpoint, hypertune) via `mecha10 init` - add your own YAML files");
        println!("there to extend the menu - each job file is a complete training config.");
        println!();
        return Ok(());
    }

    let job = select_job(&jobs, args.job.as_deref())?;
    let task_name = select_task(&ctx.working_dir, args.task.as_deref())?;

    let mut job = job.clone();
    job.task_name = Some(task_name);
    let job = &job;

    println!("Pushing configs...");
    push_configs(ctx).await?;
    println!();

    let resume_run_id_arg = args.resume_run_id.as_deref();

    run_remote(ctx, job, resume_run_id_arg).await
}

/// Prompt for which task to train on, from the project's task catalog
/// (`configs/rl/tasks/*/task.json`, see `mecha10_job_validation::list_tasks`).
fn select_task(project_root: &Path, requested: Option<&str>) -> Result<String> {
    let tasks = mecha10_job_validation::list_tasks(project_root)
        .context("failed to load task catalog from configs/rl/tasks/")?;

    if let Some(name) = requested {
        return if tasks.iter().any(|t| t == name) {
            Ok(name.to_string())
        } else {
            bail!(
                "no task named '{name}' in configs/rl/tasks/ (available: {})",
                tasks.join(", ")
            )
        };
    }

    if tasks.is_empty() {
        bail!("no tasks found under configs/rl/tasks/ - add a task.json there first.");
    }

    let selection = Select::with_theme(&ColorfulTheme::default())
        .with_prompt("Select a task")
        .items(&tasks)
        .default(0)
        .interact()
        .map_err(|e| anyhow::anyhow!("failed to get task selection: {e}"))?;

    Ok(tasks[selection].clone())
}

fn select_job<'a>(jobs: &'a [JobDefinition], requested: Option<&str>) -> Result<&'a JobDefinition> {
    if let Some(name) = requested {
        return jobs.iter().find(|j| j.name == name).with_context(|| {
            format!(
                "no job named '{name}' in configs/rl/jobs/ (available: {})",
                jobs.iter().map(|j| j.name.as_str()).collect::<Vec<_>>().join(", ")
            )
        });
    }

    let items: Vec<String> = jobs.iter().map(|j| format!("{} ({})", j.name, j.flow_type)).collect();

    let selection = Select::with_theme(&ColorfulTheme::default())
        .with_prompt("Select a training job")
        .items(&items)
        .default(0)
        .interact()
        .map_err(|e| anyhow::anyhow!("failed to get job selection: {e}"))?;

    Ok(&jobs[selection])
}

/// Response shape from `POST /api/model-deployments/jobs` - only the fields this handler prints.
#[derive(Debug, Deserialize)]
struct CreateJobResponse {
    id: String,
    name: String,
    state: String,
    mlflow_run_id: Option<String>,
    mlflow_experiment: Option<String>,
    kestra_execution_id: Option<String>,
}

#[derive(Debug, Deserialize)]
struct PriorJob {
    id: String,
    name: String,
    state: String,
    algorithm: String,
    mlflow_run_id: Option<String>,
    robot_name: Option<String>,
    created_at: String,
    hyperparameters: Option<serde_json::Value>,
}

impl PriorJob {
    fn task_name(&self) -> Option<&str> {
        self.hyperparameters.as_ref()?.get("task_name")?.as_str()
    }
}

#[derive(Debug, Deserialize)]
struct JobListResponse {
    jobs: Vec<PriorJob>,
}

#[derive(Debug, Deserialize)]
struct ArtifactEntry {
    path: String,
}

#[derive(Debug, Deserialize)]
struct JobArtifactListResponse {
    artifacts: Vec<ArtifactEntry>,
}

/// Whether `job_id`'s linked MLflow run has an actual checkpoint to resume from. Checkpoints
/// are written under the `checkpoints/` artifact path as `.zip` files.
async fn has_checkpoint(
    control_plane_url: &str,
    api_key: &str,
    client: &reqwest::Client,
    job_id: &str,
) -> Result<bool> {
    let url = format!("{control_plane_url}/api/model-deployments/jobs/{job_id}/artifacts/list?path=checkpoints");

    let response = client
        .get(&url)
        .bearer_auth(api_key)
        .send()
        .await
        .with_context(|| format!("failed to reach mecha10-api to check job {job_id}'s checkpoints"))?;

    if !response.status().is_success() {
        let status = response.status();
        let body = response.text().await.unwrap_or_default();
        bail!("failed to list job {job_id}'s artifacts ({status}): {body}");
    }

    let list: JobArtifactListResponse = response
        .json()
        .await
        .with_context(|| format!("failed to parse job {job_id}'s artifact list response"))?;

    Ok(list.artifacts.iter().any(|a| a.path.ends_with(".zip")))
}

/// List this user's prior training runs from mecha10-api, narrowed to the same
/// `algorithm`/`task_name` as the job being resumed, and let the user pick which one to
/// resume from - most recent first. Returns the picked run's `mlflow_run_id`.
async fn select_resume_run_id(
    control_plane_url: &str,
    api_key: &str,
    algorithm: &str,
    task_name: Option<&str>,
) -> Result<String> {
    let client = reqwest::Client::new();
    let url = format!("{}{}", control_plane_url, JOB_LIST_PATH);

    let response = client
        .get(&url)
        .bearer_auth(api_key)
        .send()
        .await
        .context("failed to reach mecha10-api to list prior training runs")?;

    if !response.status().is_success() {
        let status = response.status();
        let body = response.text().await.unwrap_or_default();
        bail!("failed to list prior training runs ({status}): {body}");
    }

    let list: JobListResponse = response.json().await.context("failed to parse job list response")?;

    let mut candidates: Vec<PriorJob> = list
        .jobs
        .into_iter()
        .filter(|j| j.state == "completed" && j.mlflow_run_id.is_some())
        .filter(|j| j.algorithm.eq_ignore_ascii_case(algorithm))
        .filter(|j| match task_name {
            Some(task_name) => j.task_name() == Some(task_name),
            None => true,
        })
        .collect();

    if candidates.is_empty() {
        bail!(
            "no completed '{algorithm}'{} training runs found to resume from - run a plain \
             'train' job with matching algorithm/task to completion first, or pass one directly \
             with --resume-run-id.",
            task_name.map(|t| format!(" / '{t}'")).unwrap_or_default()
        );
    }

    candidates.sort_by(|a, b| b.created_at.cmp(&a.created_at));

    println!("Checking prior runs for a resumable checkpoint...");
    let mut runs = Vec::with_capacity(candidates.len());
    for job in candidates {
        match has_checkpoint(control_plane_url, api_key, &client, &job.id).await {
            Ok(true) => runs.push(job),
            Ok(false) => {}
            Err(e) => eprintln!("warning: couldn't check job '{}' for a checkpoint: {e:#}", job.name),
        }
    }

    if runs.is_empty() {
        bail!(
            "found completed matching training runs, but none has a checkpoint artifact under \
             'checkpoints/' to resume from - pass one directly with --resume-run-id if you know \
             one has a checkpoint."
        );
    }

    let items: Vec<String> = runs
        .iter()
        .map(|j| {
            let robot = j.robot_name.as_deref().unwrap_or("-");
            format!("{} | {} | {}", j.created_at, j.name, robot)
        })
        .collect();

    let selection = Select::with_theme(&ColorfulTheme::default())
        .with_prompt("Select a prior run to resume from")
        .items(&items)
        .default(0)
        .interact()
        .map_err(|e| anyhow::anyhow!("failed to get run selection: {e}"))?;

    Ok(runs[selection].mlflow_run_id.clone().expect("filtered to Some above"))
}

/// Resolves which prior run `job` (a `resume-run` flow_type job) should resume from: the
/// `--resume-run-id` flag if given, else the interactive picker (`select_resume_run_id`). `None`
/// for any other flow_type.
async fn resolve_resume_run_id(
    job: &JobDefinition,
    resume_run_id_arg: Option<&str>,
    control_plane_url: &str,
    api_key: &str,
) -> Result<Option<String>> {
    if job.flow_type != FlowType::ResumeRun {
        return Ok(None);
    }

    if let Some(id) = resume_run_id_arg.map(str::trim).filter(|id| !id.is_empty()) {
        return Ok(Some(id.to_string()));
    }

    let algorithm = job.algorithm.clone().unwrap_or_default();
    let picked = select_resume_run_id(control_plane_url, api_key, &algorithm, job.task_name.as_deref()).await?;
    Ok(Some(picked))
}

/// Dispatch `job` remotely via mecha10-api's job-creation endpoint (Kestra).
async fn run_remote(ctx: &mut CliContext, job: &JobDefinition, resume_run_id_arg: Option<&str>) -> Result<()> {
    if !job.supports_remote() {
        bail!(
            "job '{}' uses flow_type '{}', which mecha10-api's job-creation endpoint doesn't \
             dispatch yet ('train', 'resume-run', and 'tune' are wired server-side today - see \
             mecha10-api/src/jobs/handlers.rs::train_flow_inputs/resume_flow_inputs/\
             tune_flow_inputs). Switch its flow_type to one of those.",
            job.name,
            job.flow_type
        );
    }
    if job.algorithm.as_deref().unwrap_or_default().trim().is_empty() {
        bail!(
            "job '{}' has no 'algorithm' set, which is required for remote dispatch \
             (mecha10-api's CreateJobRequest.algorithm is mandatory). Set 'algorithm:' in its YAML.",
            job.name
        );
    }

    println!("Dispatching '{}' remotely via Kestra...", job.name);
    println!();

    let credentials_service = CredentialsService::new();
    let credentials = if ctx.is_ci() {
        credentials_service
            .load()?
            .context("Not logged in. Run `mecha10 auth login` first.")?
    } else {
        mecha10_auth::ensure_authenticated_interactive(&credentials_service).await?
    };

    let project_config = ctx.load_project_config().await?;
    let control_plane_url = project_config.environments.control_plane_url();
    let url = format!("{}{}", control_plane_url, JOB_CREATE_PATH);

    let resume_run_id = resolve_resume_run_id(job, resume_run_id_arg, &control_plane_url, &credentials.api_key).await?;

    job.validate(&ctx.working_dir, resume_run_id.as_deref())?;
    let request = job.to_create_job_request(
        &ctx.working_dir,
        &project_config.robot.id,
        &project_config.name,
        resume_run_id.as_deref(),
    )?;

    let client = reqwest::Client::new();
    let response = client
        .post(&url)
        .bearer_auth(&credentials.api_key)
        .json(&request)
        .send()
        .await
        .context("failed to reach mecha10-api")?;

    if !response.status().is_success() {
        let status = response.status();
        let body = response.text().await.unwrap_or_default();
        bail!("job creation failed ({status}): {body}");
    }

    let created: CreateJobResponse = response.json().await.context("failed to parse job creation response")?;

    println!("Job created: {} ({})", created.name, created.id);
    println!("  State: {}", created.state);
    if let Some(execution_id) = &created.kestra_execution_id {
        println!("  Kestra execution: {}", execution_id);
    }
    if let Some(run_id) = &created.mlflow_run_id {
        println!("  MLflow run: {}", run_id);
    }
    if let Some(experiment) = &created.mlflow_experiment {
        println!("  MLflow experiment: {}", experiment);
    }
    println!();
    println!("Check the dashboard's ml-training job list for live progress.");
    println!();
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::services::train_jobs::FlowType;
    use std::path::PathBuf;

    fn sample_job(flow_type: FlowType) -> JobDefinition {
        JobDefinition {
            name: "train".to_string(),
            description: Some("Train from scratch".to_string()),
            flow_type,
            env_name: Some("rover".to_string()),
            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,
            n_envs: 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_trials: None,
            tune_timesteps: None,
            eval_eps: None,
            search_space: None,
            source_path: PathBuf::new(),
        }
    }

    /// Exercises `handle_train`'s pre-flight `validate_configs` call directly, confirming a
    /// broken RL job definition is caught there rather than later in `load_job_catalog`.
    #[test]
    fn train_preflight_validation_rejects_a_broken_rl_job_before_loading_the_catalog() {
        let temp = tempfile::TempDir::new().unwrap();
        std::fs::write(
            temp.path().join("mecha10.json"),
            serde_json::json!({"name": "proj", "version": "1.0.0"}).to_string(),
        )
        .unwrap();
        std::fs::create_dir_all(temp.path().join("configs/rl/jobs")).unwrap();
        std::fs::write(
            temp.path().join("configs/rl/jobs/broken.yml"),
            "not: [valid, job, definition",
        )
        .unwrap();

        let err = crate::config_validation::validate_configs(
            temp.path(),
            &[
                crate::config_validation::ConfigKind::RlJobs,
                crate::config_validation::ConfigKind::SimulationEnvironments,
                crate::config_validation::ConfigKind::ProjectAndNodes,
            ],
            None,
        )
        .unwrap_err();

        assert!(
            err.to_string().contains("RL job"),
            "expected the RL job catalog failure to surface from pre-flight validation, got: {err}"
        );
    }

    /// A project with no RL jobs, no simulation config, and a minimal valid `mecha10.json`
    /// should pass `handle_train`'s pre-flight validation set cleanly - confirming the happy
    /// path isn't blocked by kinds that simply have nothing to check.
    #[test]
    fn train_preflight_validation_passes_for_a_minimal_valid_project() {
        let temp = tempfile::TempDir::new().unwrap();
        std::fs::write(
            temp.path().join("mecha10.json"),
            serde_json::json!({"name": "proj", "version": "1.0.0"}).to_string(),
        )
        .unwrap();

        let result = crate::config_validation::validate_configs(
            temp.path(),
            &[
                crate::config_validation::ConfigKind::RlJobs,
                crate::config_validation::ConfigKind::SimulationEnvironments,
                crate::config_validation::ConfigKind::ProjectAndNodes,
            ],
            None,
        );
        assert!(result.is_ok(), "{:?}", result.err());
    }

    #[test]
    fn select_job_by_name_finds_exact_match() {
        let jobs = vec![sample_job(FlowType::Train)];
        let job = select_job(&jobs, Some("train")).unwrap();
        assert_eq!(job.name, "train");
    }

    #[test]
    fn select_job_by_name_errors_with_available_names_when_missing() {
        let jobs = vec![sample_job(FlowType::Train)];
        let err = select_job(&jobs, Some("nonexistent")).unwrap_err();
        assert!(err.to_string().contains("train"));
    }

    fn write_task(dir: &Path, name: &str) {
        let task_dir = dir.join("configs/rl/tasks").join(name);
        std::fs::create_dir_all(&task_dir).unwrap();
        std::fs::write(task_dir.join("task.json"), "{}").unwrap();
    }

    #[test]
    fn select_task_by_name_finds_exact_match() {
        let temp = tempfile::TempDir::new().unwrap();
        write_task(temp.path(), "find_target");

        let task = select_task(temp.path(), Some("find_target")).unwrap();
        assert_eq!(task, "find_target");
    }

    #[test]
    fn select_task_by_name_errors_with_available_names_when_missing() {
        let temp = tempfile::TempDir::new().unwrap();
        write_task(temp.path(), "find_target");

        let err = select_task(temp.path(), Some("nonexistent")).unwrap_err();
        assert!(err.to_string().contains("find_target"));
    }

    fn prior_job(algorithm: &str, task_name: Option<&str>, created_at: &str) -> PriorJob {
        PriorJob {
            id: "job-1".to_string(),
            name: "job-1".to_string(),
            state: "completed".to_string(),
            algorithm: algorithm.to_string(),
            mlflow_run_id: Some("run-1".to_string()),
            robot_name: None,
            created_at: created_at.to_string(),
            hyperparameters: task_name.map(|t| serde_json::json!({ "task_name": t })),
        }
    }

    #[test]
    fn prior_job_task_name_reads_it_out_of_the_hyperparameters_snapshot() {
        assert_eq!(
            prior_job("ppo", Some("navigate"), "2026-01-01T00:00:00Z").task_name(),
            Some("navigate")
        );
    }

    #[test]
    fn prior_job_task_name_is_none_without_a_hyperparameters_snapshot() {
        assert_eq!(prior_job("ppo", None, "2026-01-01T00:00:00Z").task_name(), None);
    }

    /// Mirrors `select_resume_run_id`'s own filter/sort - not a call through that fn (it's an
    /// interactive HTTP handler), but the same algorithm/task-matching and recency-ordering logic
    /// this test exercises directly against `PriorJob`.
    #[test]
    fn resume_candidates_are_filtered_to_matching_algorithm_and_task_then_sorted_most_recent_first() {
        let jobs = [
            prior_job("ppo", Some("navigate"), "2026-01-01T00:00:00Z"),
            prior_job("sac", Some("navigate"), "2026-01-03T00:00:00Z"),
            prior_job("ppo", Some("find_target"), "2026-01-02T00:00:00Z"),
            prior_job("ppo", Some("navigate"), "2026-01-05T00:00:00Z"),
        ];

        let algorithm = "ppo";
        let task_name = Some("navigate");
        let mut candidates: Vec<&PriorJob> = jobs
            .iter()
            .filter(|j| j.algorithm.eq_ignore_ascii_case(algorithm))
            .filter(|j| match task_name {
                Some(task_name) => j.task_name() == Some(task_name),
                None => true,
            })
            .collect();
        candidates.sort_by(|a, b| b.created_at.cmp(&a.created_at));

        let created_ats: Vec<&str> = candidates.iter().map(|j| j.created_at.as_str()).collect();
        assert_eq!(created_ats, vec!["2026-01-05T00:00:00Z", "2026-01-01T00:00:00Z"]);
    }
}