mecha10-cli 0.6.3

//! `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, JobDefinition};
use anyhow::{bail, Context, Result};
use dialoguer::{theme::ColorfulTheme, Select};
use serde::Deserialize;
use std::path::Path;
use std::process::Stdio;

/// 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. See configs/rl/local.yml for the underlying 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<()> {
    println!("Running '{}' locally...", job.name);
    println!();

    let credentials = CredentialsService::new().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.",
    )?;

    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(&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>,
}

/// 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 (only the plain 'train' flow is wired server-side today - see \
             mecha10-api/src/jobs/handlers.rs::train_flow_inputs). Run it locally instead, or \
             add a plain 'train' job definition.",
            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 = CredentialsService::new()
        .load()?
        .context("Not logged in. Run `mecha10 auth login` first.")?;

    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 request = job.to_create_job_request(&project_config.robot.id);

    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,
            config: "local".to_string(),
            env_name: Some("rover".to_string()),
            algorithm: Some("ppo".to_string()),
            timesteps: None,
            n_steps: None,
            epochs: None,
            lr: 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();
        // No .venv under `temp`, so this exercises the PATH-probing fallback. This environment
        // is assumed to have some Python 3 available (same assumption `sim::find_python` makes).
        let python = find_project_python(temp.path());
        assert!(python.is_ok());
    }
}