mecha10-cli 0.6.5

//! `mecha10 train` command handler (LAB-1776).
//!
//! Interactive flow: menu of job definitions (bundled + project-defined, `configs/rl/jobs/*.yml`
//! - see `services::train_jobs`) -> pick one -> pick local or remote dispatch.
//!
//! - **Local**: spawns `python3 -m rl_task_runner.run <train|tune>` in-process/subprocess against
//!   the project's own `configs/rl/`, with the metrics relay tracker enabled (`RelayTracker`,
//!   `packages/rl-task-runner/rl_task_runner/services/relay_tracker.py`) - no Kestra/MLflow
//!   credentials needed on this machine.
//! - **Remote**: `POST`s to mecha10-api's job-creation endpoint
//!   (`/api/model-deployments/jobs`), the exact same contract `packages/dashboard`'s job-launch
//!   flow already uses (PR #278/LAB-1431) - see `services::train_jobs::CreateJobRequest`.

use crate::commands::TrainArgs;
use crate::context::CliContext;
use crate::services::credentials::CredentialsService;
use crate::services::train_jobs::{load_job_catalog, FlowType, JobDefinition};
use anyhow::{bail, Context, Result};
use dialoguer::{theme::ColorfulTheme, Select};
use serde::Deserialize;
use std::path::Path;
use std::process::Stdio;

/// Jobs-list endpoint, used to build the "which run should this resume from" picker -
/// see `select_resume_run_id`.
const JOB_LIST_PATH: &str = "/api/model-deployments/jobs";

/// Job-creation endpoint, appended to the project's control plane URL. Matches the exact
/// contract `packages/dashboard`'s `lib/api/ml-training.ts` already uses against
/// `mecha10-api`'s `jobs::router` - do not reinvent this shape (LAB-1776 requirement).
const JOB_CREATE_PATH: &str = "/api/model-deployments/jobs";

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum DispatchTarget {
    Local,
    Remote,
}

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

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

    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 target = select_target(args.target.as_deref())?;

    match target {
        DispatchTarget::Local => run_local(ctx, job).await,
        DispatchTarget::Remote => run_remote(ctx, job).await,
    }
}

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| match &j.description {
            Some(desc) => format!("{} ({}) - {}", j.name, j.flow_type, desc),
            None => 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])
}

fn select_target(requested: Option<&str>) -> Result<DispatchTarget> {
    if let Some(target) = requested {
        return match target.to_lowercase().as_str() {
            "local" => Ok(DispatchTarget::Local),
            "remote" => Ok(DispatchTarget::Remote),
            other => bail!("unknown --target '{other}' (expected 'local' or 'remote')"),
        };
    }

    let items = ["Local (run on this machine)", "Remote (dispatch via Kestra/SaaS)"];
    let selection = Select::with_theme(&ColorfulTheme::default())
        .with_prompt("Where should this job run?")
        .items(items)
        .default(0)
        .interact()
        .map_err(|e| anyhow::anyhow!("failed to get dispatch target selection: {e}"))?;

    Ok(if selection == 0 {
        DispatchTarget::Local
    } else {
        DispatchTarget::Remote
    })
}

/// Run `job` locally: `python3 -m rl_task_runner.run <mode>`, inheriting stdio so the Python
/// side's own logging (including relay-report failures) streams straight to the user, mirroring
/// `sim::run_mujoco_cli`'s streaming-subprocess pattern.
async fn run_local(ctx: &mut CliContext, job: &JobDefinition) -> Result<()> {
    job.validate(&ctx.working_dir)?;

    println!("Running '{}' locally...", 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 - local training still reports \
             metrics through mecha10-api's relay, which needs an API key.",
        )?
    } 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 python = find_project_python(&ctx.working_dir)?;
    let log_level = ctx.log_level.to_string().to_lowercase();
    let env_vars = job.local_env_vars(&ctx.working_dir, &control_plane_url, &credentials.api_key, &log_level);
    let mode = job.run_mode();

    let mut command = tokio::process::Command::new(&python);
    command
        .args(["-m", "rl_task_runner.run", mode])
        .current_dir(&ctx.working_dir)
        .stdin(Stdio::inherit())
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit());
    for (key, value) in &env_vars {
        command.env(key, value);
    }

    let status = command
        .status()
        .await
        .with_context(|| format!("failed to spawn '{python} -m rl_task_runner.run {mode}'"))?;

    if !status.success() {
        bail!(
            "training job '{}' exited with {}",
            job.name,
            status
                .code()
                .map_or_else(|| "a signal".to_string(), |code| code.to_string())
        );
    }

    println!();
    println!("Training job '{}' finished.", job.name);
    Ok(())
}

/// Resolve a Python interpreter for local training: prefer the project's own `.venv` (created by
/// `mecha10 setup`, which installs `rl-task-runner` from `requirements.txt`), falling back to a
/// `python3`/`python` found on `PATH`.
fn find_project_python(project_root: &Path) -> Result<String> {
    let venv_python = if cfg!(windows) {
        project_root.join(".venv").join("Scripts").join("python.exe")
    } else {
        project_root.join(".venv").join("bin").join("python")
    };
    if venv_python.exists() {
        return Ok(venv_python.display().to_string());
    }

    for candidate in ["python3", "python"] {
        if let Ok(output) = std::process::Command::new(candidate).arg("--version").output() {
            let version = format!(
                "{}{}",
                String::from_utf8_lossy(&output.stdout),
                String::from_utf8_lossy(&output.stderr)
            );
            if output.status.success() && version.contains("Python 3") {
                return Ok(candidate.to_string());
            }
        }
    }

    bail!(
        "Python 3 not found (checked {} and PATH). Run `mecha10 setup` to create a project venv \
         with rl-task-runner installed.",
        venv_python.display()
    )
}

/// Response shape from `POST /api/model-deployments/jobs` - only the fields this handler prints;
/// see `mecha10-api/src/jobs/types.rs::CreateJobResponse` for the full/authoritative shape.
#[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>,
}

/// One entry from `GET /api/model-deployments/jobs` - only the fields the resume-run picker
/// needs to display, to check for a checkpoint, to filter to the same algorithm/task as the
/// job being resumed, and to forward as `resume_run_id`; see
/// `mecha10-api/src/jobs/types.rs::Job` for the full/authoritative shape.
#[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,
    /// `Job::hyperparameters`'s creation-time snapshot - only `task_name` is read out of it
    /// here (`Job` has no top-level `task_name` field; it only lives inside this blob).
    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>,
}

/// One entry from `GET /api/model-deployments/jobs/:id/artifacts/list` - only the field
/// `has_checkpoint` needs; see `mecha10-api/src/jobs/mlflow_client.rs::ArtifactEntry` for the
/// full shape.
#[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. Mirrors
/// `rl_task_runner.run._resolve_warm_start`'s own lookup exactly (checkpoints are only ever
/// written under the `checkpoints/` artifact path, as `.zip` files -
/// `services.callbacks.MLflowCheckpointCallback`) - a `completed` job that finished before its
/// first checkpoint interval has no `.zip` there yet, and would otherwise only fail once deep
/// inside the trainer container (`run.py` raises when `_resolve_warm_start` finds none).
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 (a checkpoint from a different
/// algorithm or task can't be warm-started into this run - `rl_task_runner.run`'s
/// `warm_start` assumes a matching model architecture), and let the user pick which one to
/// resume from - most recent first. Returns the picked run's `mlflow_run_id`, the value
/// `resume-run.yml` forwards to the trainer as `RESUME_RUN_ID`, which `rl_task_runner.run`
/// resolves (via the MLflow query API) into the run's last checkpoint artifact - that same
/// run also anchors `resume_flow_inputs`' `model_name`, so the checkpoint's own prior job
/// doubles as the reference a resumed run continues (both its weights and its registry
/// model lineage), with no separate job-id field needed.
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")?;

    // Only a `completed` run has a finished checkpoint worth resuming from - a `running`/
    // `pending` run's checkpoint is still being written, and `failed`/`cancelled` runs aren't
    // guaranteed to have produced one at all. Also require the same algorithm/task as the job
    // being resumed - see this fn's doc comment.
    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 paste a run ID \
             directly into this job's 'resume_run_id:' field.",
            task_name.map(|t| format!(" / '{t}'")).unwrap_or_default()
        );
    }

    // mecha10-api already returns jobs most-recent-first (`ORDER BY jobs.created_at DESC`) -
    // sort again here defensively so the picker is correct even if that changes, since "most
    // recent first" is this picker's whole point.
    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 - paste a run ID directly into this job's \
             'resume_run_id:' field 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"))
}

/// Dispatch `job` remotely via mecha10-api's job-creation endpoint (Kestra).
async fn run_remote(ctx: &mut CliContext, job: &JobDefinition) -> Result<()> {
    if !job.supports_remote() {
        bail!(
            "job '{}' uses flow_type '{}', which mecha10-api's job-creation endpoint doesn't \
             dispatch yet ('train' and 'resume-run' are wired server-side today - see \
             mecha10-api/src/jobs/handlers.rs::train_flow_inputs/resume_flow_inputs). Run it \
             locally instead, or 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, or run it locally, where the config file's own algorithm_name is used instead.",
            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);

    // `train-from-checkpoint`-style jobs (flow_type: resume-run) need a target run to warm-start
    // from - if the job's own YAML didn't already pin one, ask which prior run to resume,
    // listing this user's own jobs (LAB per user request: "list prior runs in the cli, let the
    // user pick, inject the run id").
    let mut job = job.clone();
    if job.flow_type == FlowType::ResumeRun && job.resume_run_id.as_deref().unwrap_or_default().trim().is_empty() {
        let algorithm = job.algorithm.clone().unwrap_or_default();
        job.resume_run_id = Some(
            select_resume_run_id(
                &control_plane_url,
                &credentials.api_key,
                &algorithm,
                job.task_name.as_deref(),
            )
            // (picker fills `resume_run_id`; the full field/task-catalog validation below runs
            // after this so it covers the picked value too)
            .await?,
        );
    }

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

    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,
            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,
            resume_run_id: None,
            tune_run_id: None,
            n_trials: None,
            tune_timesteps: None,
            eval_eps: None,
            source_path: PathBuf::new(),
        }
    }

    #[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"));
    }

    #[test]
    fn select_target_parses_local_and_remote_case_insensitively() {
        assert_eq!(select_target(Some("local")).unwrap(), DispatchTarget::Local);
        assert_eq!(select_target(Some("REMOTE")).unwrap(), DispatchTarget::Remote);
    }

    #[test]
    fn select_target_rejects_unknown_value() {
        let err = select_target(Some("saas")).unwrap_err();
        assert!(err.to_string().contains("saas"));
    }

    #[test]
    fn find_project_python_falls_back_to_path_when_no_venv() {
        let temp = tempfile::TempDir::new().unwrap();
        let python = find_project_python(temp.path());
        assert!(python.is_ok());
    }

    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"]);
    }
}