//! K8s Job Dispatcher
//!
//! Creates `batch/v1` Jobs on the gpu-training pool for on-demand model training.
//! Each Job runs a main training container with GPU access and a native sidecar
//! (K8s 1.28+ `restartPolicy: Always` init container) that uploads artifacts to S3.

use std::collections::BTreeMap;

use anyhow::{Context, Result};
use k8s_openapi::api::batch::v1::{Job, JobSpec};
use k8s_openapi::api::core::v1::{
    Container, EmptyDirVolumeSource, EnvFromSource, EnvVar,
    LocalObjectReference, PersistentVolumeClaimVolumeSource, PodSpec, PodTemplateSpec,
    ResourceRequirements, SecretEnvSource, SecurityContext, Toleration, Volume, VolumeMount,
};
use k8s_openapi::apimachinery::pkg::api::resource::Quantity;
use k8s_openapi::apimachinery::pkg::apis::meta::v1::ObjectMeta;
use kube::api::{DeleteParams, PostParams};
use kube::{Api, Client};
use tracing::info;
use uuid::Uuid;

// ---------------------------------------------------------------------------
// Configuration
// ---------------------------------------------------------------------------

/// Configuration for the K8s job dispatcher.
///
/// All fields can be populated from environment variables via [`DispatcherConfig::from_env`].
#[derive(Debug, Clone)]
pub struct DispatcherConfig {
    pub namespace: String,
    pub runtime_image: String,
    pub data_pvc: String,
    pub s3_bucket: String,
    pub s3_endpoint: String,
    pub s3_region: String,
    pub binaries_bucket: String,
    pub callback_endpoint: String,
}

impl Default for DispatcherConfig {
    fn default() -> Self {
        Self::from_env()
    }
}

impl DispatcherConfig {
    /// Reads configuration from environment variables.
    ///
    /// Required env vars: S3_ENDPOINT, TRAINING_RUNTIME_IMAGE, CALLBACK_ENDPOINT.
    /// Falls back to defaults only for non-URL values (namespace, bucket names, PVC names).
    #[allow(clippy::expect_used)] // Intentional: fail fast at service startup if required env vars are missing
    pub fn from_env() -> Self {
        Self {
            namespace: std::env::var("K8S_NAMESPACE").unwrap_or_else(|_| "foxhunt".to_string()),
            runtime_image: std::env::var("TRAINING_RUNTIME_IMAGE")
                .expect("TRAINING_RUNTIME_IMAGE env var required"),
            data_pvc: std::env::var("TRAINING_DATA_PVC")
                .unwrap_or_else(|_| "training-data-pvc".to_string()),
            s3_bucket: std::env::var("S3_BUCKET")
                .unwrap_or_else(|_| "foxhunt-models".to_string()),
            s3_endpoint: std::env::var("S3_ENDPOINT")
                .expect("S3_ENDPOINT env var required"),
            s3_region: std::env::var("S3_REGION")
                .unwrap_or_else(|_| "fr-par".to_string()),
            binaries_bucket: std::env::var("BINARIES_BUCKET")
                .unwrap_or_else(|_| "foxhunt-binaries".to_string()),
            callback_endpoint: std::env::var("CALLBACK_ENDPOINT")
                .expect("CALLBACK_ENDPOINT env var required"),
        }
    }
}

// ---------------------------------------------------------------------------
// Job parameters
// ---------------------------------------------------------------------------

/// Parameters for a training job dispatch.
#[derive(Debug, Clone)]
pub struct TrainingJobParams {
    pub job_id: Uuid,
    pub model_type: String,
    pub symbol: String,
    pub epochs: u32,
    pub binary: String,
    /// Additional CLI arguments appended after --symbol and --data-dir.
    /// Used for hyperopt-specific flags (--trials, --model, --parallel, --output).
    pub extra_args: Vec<String>,
}

/// Returns the correct training binary name for a given model type.
///
/// RL models (DQN, PPO) use `train_baseline_rl`; all supervised models use
/// `train_baseline_supervised`.
pub fn training_binary_for_model(model_type: &str) -> &'static str {
    match model_type.to_lowercase().as_str() {
        "dqn" | "ppo" => "train_baseline_rl",
        _ => "train_baseline_supervised",
    }
}

/// Returns the correct hyperopt binary name for a given model type.
///
/// RL models (DQN, PPO) use `hyperopt_baseline_rl`; all supervised models use
/// `hyperopt_baseline_supervised`.
pub fn hyperopt_binary_for_model(model_type: &str) -> &'static str {
    match model_type.to_lowercase().as_str() {
        "dqn" | "ppo" => "hyperopt_baseline_rl",
        _ => "hyperopt_baseline_supervised",
    }
}

// ---------------------------------------------------------------------------
// Dispatcher
// ---------------------------------------------------------------------------

/// Dispatches K8s `batch/v1` Jobs for GPU model training.
pub struct K8sDispatcher {
    client: Client,
    config: DispatcherConfig,
}

impl K8sDispatcher {
    /// Creates a new dispatcher using in-cluster Kubernetes configuration.
    pub async fn new(config: DispatcherConfig) -> Result<Self> {
        let client =
            Client::try_default().await.context("failed to create K8s client from in-cluster or kubeconfig")?;
        Ok(Self { client, config })
    }

    /// Submits a training Job to the cluster and returns its K8s name.
    pub async fn dispatch(&self, params: &TrainingJobParams) -> Result<String> {
        let job = build_job_spec(&self.config, params);
        let job_name = job
            .metadata
            .name
            .clone()
            .unwrap_or_else(|| "unknown".to_string());

        let jobs: Api<Job> = Api::namespaced(self.client.clone(), &self.config.namespace);

        info!(
            job_name = %job_name,
            model_type = %params.model_type,
            symbol = %params.symbol,
            "dispatching training job"
        );

        jobs.create(&PostParams::default(), &job)
            .await
            .with_context(|| format!("failed to create K8s Job {job_name}"))?;

        info!(job_name = %job_name, "training job created");
        Ok(job_name)
    }

    /// Returns the current phase of a training Job (`Active`, `Complete`, `Failed`),
    /// or `None` if the Job does not exist.
    pub async fn get_job_status(&self, job_name: &str) -> Result<Option<String>> {
        let jobs: Api<Job> = Api::namespaced(self.client.clone(), &self.config.namespace);

        match jobs.get_opt(job_name).await.context("failed to query job status")? {
            Some(job) => {
                let status = job.status.as_ref();
                if status.and_then(|s| s.succeeded).unwrap_or(0) > 0 {
                    Ok(Some("Complete".to_string()))
                } else if status.and_then(|s| s.failed).unwrap_or(0) > 0 {
                    Ok(Some("Failed".to_string()))
                } else if status.and_then(|s| s.active).unwrap_or(0) > 0 {
                    Ok(Some("Active".to_string()))
                } else {
                    Ok(Some("Pending".to_string()))
                }
            }
            None => Ok(None),
        }
    }

    /// Deletes a completed or failed training Job from the cluster.
    pub async fn delete_job(&self, job_name: &str) -> Result<()> {
        let jobs: Api<Job> = Api::namespaced(self.client.clone(), &self.config.namespace);

        info!(job_name = %job_name, "deleting training job");

        let dp = DeleteParams {
            propagation_policy: Some(kube::api::PropagationPolicy::Background),
            ..Default::default()
        };

        jobs.delete(job_name, &dp)
            .await
            .with_context(|| format!("failed to delete K8s Job {job_name}"))?;

        Ok(())
    }
}

// ---------------------------------------------------------------------------
// Job manifest builder (pure function — no cluster access)
// ---------------------------------------------------------------------------

/// Builds a complete `batch/v1` Job manifest without touching the cluster.
///
/// Matches the production `infra/k8s/training/job-template.yaml`:
/// - fetch-binaries initContainer (rclone from S3)
/// - uploader native sidecar (K8s 1.28+)
/// - main training container with GPU
/// - emptyDir for output and binaries, PVC for training data
pub fn build_job_spec(config: &DispatcherConfig, params: &TrainingJobParams) -> Job {
    let raw_name = format!("training-{}-{}", params.model_type, params.job_id);
    // K8s names are limited to 63 characters and must be lowercase DNS-compatible.
    let job_name = sanitize_k8s_name(&raw_name);

    let mut labels = BTreeMap::new();
    labels.insert(
        "app.kubernetes.io/name".to_string(),
        "training".to_string(),
    );
    labels.insert(
        "foxhunt/job-type".to_string(),
        if params.binary.contains("hyperopt") { "hyperopt" } else { "training" }.to_string(),
    );
    labels.insert(
        "foxhunt/model".to_string(),
        params.model_type.to_lowercase(),
    );
    labels.insert(
        "foxhunt/job-id".to_string(),
        params.job_id.to_string(),
    );

    // -- fetch-binaries initContainer (runs first, downloads from S3) ---------
    let fetch_binaries = Container {
        name: "fetch-binaries".to_string(),
        image: Some(config.runtime_image.clone()),
        security_context: Some(SecurityContext {
            run_as_user: Some(0),
            ..Default::default()
        }),
        command: Some(vec!["/bin/sh".to_string(), "-c".to_string()]),
        args: Some(vec![format!(
            "set -e\n\
             export RCLONE_S3_PROVIDER=Minio RCLONE_S3_ENDPOINT=$S3_ENDPOINT RCLONE_S3_REGION=$S3_REGION\n\
             rclone sync :s3:{}/latest/training/ /binaries/\n\
             chmod +x /binaries/*\n\
             ls -lh /binaries/\n\
             echo \"Fetched training binaries from S3\"",
            config.binaries_bucket
        )]),
        env: Some(vec![
            env_var("S3_ENDPOINT", &config.s3_endpoint),
            env_var("S3_REGION", &config.s3_region),
            EnvVar {
                name: "RCLONE_S3_ACCESS_KEY_ID".to_string(),
                value_from: Some(k8s_openapi::api::core::v1::EnvVarSource {
                    secret_key_ref: Some(k8s_openapi::api::core::v1::SecretKeySelector {
                        name: "s3-credentials".to_string(),
                        key: "access-key".to_string(),
                        optional: Some(false),
                    }),
                    ..Default::default()
                }),
                ..Default::default()
            },
            EnvVar {
                name: "RCLONE_S3_SECRET_ACCESS_KEY".to_string(),
                value_from: Some(k8s_openapi::api::core::v1::EnvVarSource {
                    secret_key_ref: Some(k8s_openapi::api::core::v1::SecretKeySelector {
                        name: "s3-credentials".to_string(),
                        key: "secret-key".to_string(),
                        optional: Some(false),
                    }),
                    ..Default::default()
                }),
                ..Default::default()
            },
        ]),
        volume_mounts: Some(vec![
            VolumeMount {
                name: "binaries".to_string(),
                mount_path: "/binaries".to_string(),
                ..Default::default()
            },
        ]),
        resources: Some(ResourceRequirements {
            requests: Some(BTreeMap::from([
                ("cpu".to_string(), Quantity("100m".to_string())),
                ("memory".to_string(), Quantity("64Mi".to_string())),
            ])),
            limits: Some(BTreeMap::from([
                ("cpu".to_string(), Quantity("500m".to_string())),
                ("memory".to_string(), Quantity("256Mi".to_string())),
            ])),
            ..Default::default()
        }),
        ..Default::default()
    };

    // -- Uploader sidecar (native K8s 1.28+) ----------------------------------
    let uploader = Container {
        name: "uploader".to_string(),
        image: Some(config.runtime_image.clone()),
        restart_policy: Some("Always".to_string()), // native sidecar
        command: Some(vec!["/binaries/training_uploader".to_string()]),
        env: Some(vec![
            env_var("JOB_ID", &params.job_id.to_string()),
            env_var("S3_BUCKET", &config.s3_bucket),
            env_var("S3_ENDPOINT", &config.s3_endpoint),
            env_var("S3_REGION", &config.s3_region),
            env_var("S3_PREFIX", &format!("models/{}", params.job_id)),
            env_var("CALLBACK_ENDPOINT", &config.callback_endpoint),
            env_var("RUST_LOG", "info"),
        ]),
        env_from: Some(vec![EnvFromSource {
            secret_ref: Some(SecretEnvSource {
                name: "s3-credentials".to_string(),
                optional: Some(false),
            }),
            ..Default::default()
        }]),
        volume_mounts: Some(vec![
            VolumeMount {
                name: "output".to_string(),
                mount_path: "/output".to_string(),
                read_only: Some(true),
                ..Default::default()
            },
            VolumeMount {
                name: "binaries".to_string(),
                mount_path: "/binaries".to_string(),
                read_only: Some(true),
                ..Default::default()
            },
        ]),
        resources: Some(ResourceRequirements {
            requests: Some(BTreeMap::from([
                ("cpu".to_string(), Quantity("100m".to_string())),
                ("memory".to_string(), Quantity("128Mi".to_string())),
            ])),
            limits: Some(BTreeMap::from([
                ("cpu".to_string(), Quantity("500m".to_string())),
                ("memory".to_string(), Quantity("256Mi".to_string())),
            ])),
            ..Default::default()
        }),
        ..Default::default()
    };

    // -- Main training container ----------------------------------------------
    let training = Container {
        name: "training".to_string(),
        image: Some(config.runtime_image.clone()),
        command: Some(vec![format!("/binaries/{}", params.binary)]),
        args: Some({
            let mut args = vec![
                format!("--symbol={}", params.symbol),
                "--data-dir=/data/futures-baseline".to_string(),
            ];
            if params.extra_args.is_empty() {
                // Default training args when no extra_args provided
                args.push("--output-dir=/output".to_string());
                args.push(format!("--epochs={}", params.epochs));
            } else {
                args.extend(params.extra_args.iter().cloned());
            }
            args
        }),
        env: Some(vec![
            env_var("RUST_LOG", "info"),
            env_var("SQLX_OFFLINE", "true"),
        ]),
        volume_mounts: Some(vec![
            VolumeMount {
                name: "training-data".to_string(),
                mount_path: "/data".to_string(),
                read_only: Some(true),
                ..Default::default()
            },
            VolumeMount {
                name: "output".to_string(),
                mount_path: "/output".to_string(),
                ..Default::default()
            },
            VolumeMount {
                name: "binaries".to_string(),
                mount_path: "/binaries".to_string(),
                read_only: Some(true),
                ..Default::default()
            },
        ]),
        resources: Some(ResourceRequirements {
            requests: Some(BTreeMap::from([
                ("nvidia.com/gpu".to_string(), Quantity("1".to_string())),
                ("cpu".to_string(), Quantity("4".to_string())),
                ("memory".to_string(), Quantity("16Gi".to_string())),
            ])),
            limits: Some(BTreeMap::from([
                ("nvidia.com/gpu".to_string(), Quantity("1".to_string())),
                ("cpu".to_string(), Quantity("8".to_string())),
                ("memory".to_string(), Quantity("32Gi".to_string())),
            ])),
            ..Default::default()
        }),
        ..Default::default()
    };

    // -- Volumes ---------------------------------------------------------------
    let volumes = vec![
        Volume {
            name: "training-data".to_string(),
            persistent_volume_claim: Some(PersistentVolumeClaimVolumeSource {
                claim_name: config.data_pvc.clone(),
                read_only: Some(true),
            }),
            ..Default::default()
        },
        Volume {
            name: "output".to_string(),
            empty_dir: Some(EmptyDirVolumeSource {
                size_limit: Some(Quantity("2Gi".to_string())),
                ..Default::default()
            }),
            ..Default::default()
        },
        Volume {
            name: "binaries".to_string(),
            empty_dir: Some(EmptyDirVolumeSource {
                size_limit: Some(Quantity("500Mi".to_string())),
                ..Default::default()
            }),
            ..Default::default()
        },
    ];

    // -- Node selector ---------------------------------------------------------
    let mut node_selector = BTreeMap::new();
    node_selector.insert(
        "k8s.scaleway.com/pool-name".to_string(),
        "gpu-training".to_string(),
    );

    // -- Pod spec --------------------------------------------------------------
    let pod_spec = PodSpec {
        // fetch-binaries runs first (sequential), then uploader starts as native sidecar
        init_containers: Some(vec![fetch_binaries, uploader]),
        containers: vec![training],
        volumes: Some(volumes),
        node_selector: Some(node_selector),
        tolerations: Some(vec![
            Toleration {
                key: Some("nvidia.com/gpu".to_string()),
                operator: Some("Exists".to_string()),
                effect: Some("NoSchedule".to_string()),
                ..Default::default()
            },
            // Cilium CNI takes ~30s on fresh scale-from-zero nodes
            Toleration {
                key: Some("node.cilium.io/agent-not-ready".to_string()),
                operator: Some("Exists".to_string()),
                effect: Some("NoSchedule".to_string()),
                ..Default::default()
            },
        ]),
        image_pull_secrets: Some(vec![LocalObjectReference {
            name: "gitlab-registry".to_string(),
        }]),
        restart_policy: Some("Never".to_string()),
        ..Default::default()
    };

    // -- Job -------------------------------------------------------------------
    Job {
        metadata: ObjectMeta {
            name: Some(job_name),
            namespace: Some(config.namespace.clone()),
            labels: Some(labels.clone()),
            ..Default::default()
        },
        spec: Some(JobSpec {
            backoff_limit: Some(1),
            active_deadline_seconds: Some(21600), // 6 hours
            ttl_seconds_after_finished: Some(600),
            template: PodTemplateSpec {
                metadata: Some(ObjectMeta {
                    labels: Some(labels),
                    ..Default::default()
                }),
                spec: Some(pod_spec),
            },
            ..Default::default()
        }),
        ..Default::default()
    }
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// Creates a simple `EnvVar` with a name and value.
fn env_var(name: &str, value: &str) -> EnvVar {
    EnvVar {
        name: name.to_string(),
        value: Some(value.to_string()),
        ..Default::default()
    }
}

/// Sanitizes a string into a valid K8s resource name (lowercase, max 63 chars,
/// DNS-1123 label compliant).
fn sanitize_k8s_name(name: &str) -> String {
    let lower = name.to_lowercase();
    let sanitized: String = lower
        .chars()
        .map(|c| if c.is_ascii_alphanumeric() || c == '-' { c } else { '-' })
        .collect();
    // Trim leading/trailing hyphens and truncate to 63 chars.
    let trimmed = sanitized.trim_matches('-');
    let truncated = trimmed.get(..63).unwrap_or(trimmed);
    truncated.trim_end_matches('-').to_string()
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;

    fn test_config() -> DispatcherConfig {
        DispatcherConfig {
            namespace: "foxhunt-test".to_string(),
            runtime_image: "registry.example.com/foxhunt-training-runtime:v1".to_string(),
            data_pvc: "training-data-pvc".to_string(),
            s3_bucket: "test-bucket".to_string(),
            s3_endpoint: "https://s3.test.example.com".to_string(),
            s3_region: "us-test-1".to_string(),
            binaries_bucket: "test-binaries".to_string(),
            callback_endpoint: "http://callback.test.svc:50053".to_string(),
        }
    }

    fn test_params() -> TrainingJobParams {
        TrainingJobParams {
            job_id: Uuid::parse_str("550e8400-e29b-41d4-a716-446655440000").unwrap(),
            model_type: "tft".to_string(),
            symbol: "ES.FUT".to_string(),
            epochs: 50,
            binary: "train_baseline_supervised".to_string(),
            extra_args: vec![],
        }
    }

    #[test]
    fn test_training_binary_for_model() {
        assert_eq!(training_binary_for_model("dqn"), "train_baseline_rl");
        assert_eq!(training_binary_for_model("DQN"), "train_baseline_rl");
        assert_eq!(training_binary_for_model("ppo"), "train_baseline_rl");
        assert_eq!(training_binary_for_model("PPO"), "train_baseline_rl");

        assert_eq!(training_binary_for_model("tft"), "train_baseline_supervised");
        assert_eq!(training_binary_for_model("mamba2"), "train_baseline_supervised");
        assert_eq!(training_binary_for_model("tggn"), "train_baseline_supervised");
        assert_eq!(training_binary_for_model("tlob"), "train_baseline_supervised");
        assert_eq!(training_binary_for_model("liquid"), "train_baseline_supervised");
        assert_eq!(training_binary_for_model("kan"), "train_baseline_supervised");
        assert_eq!(training_binary_for_model("xlstm"), "train_baseline_supervised");
        assert_eq!(
            training_binary_for_model("diffusion"),
            "train_baseline_supervised"
        );
    }

    #[test]
    fn test_build_job_spec() {
        let config = test_config();
        let params = test_params();

        let job = build_job_spec(&config, &params);

        // -- Metadata --
        let meta = &job.metadata;
        let name = meta.name.as_deref().unwrap();
        assert!(
            name.starts_with("training-tft-"),
            "job name should start with training-tft-, got: {name}"
        );
        assert!(name.len() <= 63, "job name must be <= 63 chars");

        let labels = meta.labels.as_ref().unwrap();
        assert_eq!(labels.get("app.kubernetes.io/name").unwrap(), "training");
        assert_eq!(labels.get("foxhunt/job-type").unwrap(), "training");
        assert_eq!(labels.get("foxhunt/model").unwrap(), "tft");
        assert_eq!(
            labels.get("foxhunt/job-id").unwrap(),
            "550e8400-e29b-41d4-a716-446655440000"
        );

        assert_eq!(meta.namespace.as_deref().unwrap(), "foxhunt-test");

        // -- Job spec --
        let spec = job.spec.as_ref().unwrap();
        assert_eq!(spec.backoff_limit, Some(1));
        assert_eq!(spec.active_deadline_seconds, Some(21600)); // 6 hours
        assert_eq!(spec.ttl_seconds_after_finished, Some(600));

        // -- Pod spec --
        let pod_spec = spec.template.spec.as_ref().unwrap();

        // Node selector
        let ns = pod_spec.node_selector.as_ref().unwrap();
        assert_eq!(
            ns.get("k8s.scaleway.com/pool-name").unwrap(),
            "gpu-training"
        );

        // Tolerations (GPU + Cilium)
        let tols = pod_spec.tolerations.as_ref().unwrap();
        assert_eq!(tols.len(), 2);
        assert_eq!(tols[0].key.as_deref().unwrap(), "nvidia.com/gpu");
        assert_eq!(tols[1].key.as_deref().unwrap(), "node.cilium.io/agent-not-ready");

        // Image pull secrets
        let ips = pod_spec.image_pull_secrets.as_ref().unwrap();
        assert_eq!(ips[0].name, "gitlab-registry");

        // Restart policy
        assert_eq!(pod_spec.restart_policy.as_deref().unwrap(), "Never");

        // -- Init containers: fetch-binaries + uploader sidecar --
        let init_containers = pod_spec.init_containers.as_ref().unwrap();
        assert_eq!(init_containers.len(), 2);

        // fetch-binaries (sequential init, runs first)
        let fetcher = &init_containers[0];
        assert_eq!(fetcher.name, "fetch-binaries");
        assert!(fetcher.restart_policy.is_none(), "fetcher is not a sidecar");
        assert_eq!(
            fetcher.security_context.as_ref().unwrap().run_as_user,
            Some(0)
        );

        // uploader (native sidecar)
        let sidecar = &init_containers[1];
        assert_eq!(sidecar.name, "uploader");
        assert_eq!(
            sidecar.restart_policy.as_deref().unwrap(),
            "Always",
            "sidecar must have restart_policy Always for native K8s 1.28+ sidecar"
        );
        assert_eq!(
            sidecar.image.as_deref().unwrap(),
            "registry.example.com/foxhunt-training-runtime:v1"
        );

        // Sidecar env
        let sidecar_env = sidecar.env.as_ref().unwrap();
        let env_names: Vec<&str> = sidecar_env.iter().map(|e| e.name.as_str()).collect();
        assert!(env_names.contains(&"JOB_ID"));
        assert!(env_names.contains(&"S3_BUCKET"));
        assert!(env_names.contains(&"S3_ENDPOINT"));
        assert!(env_names.contains(&"S3_PREFIX"));
        assert!(env_names.contains(&"CALLBACK_ENDPOINT"));

        // Sidecar volume mounts (output + binaries)
        let sidecar_mounts = sidecar.volume_mounts.as_ref().unwrap();
        assert_eq!(sidecar_mounts.len(), 2);
        assert!(sidecar_mounts.iter().any(|m| m.name == "output"));
        assert!(sidecar_mounts.iter().any(|m| m.name == "binaries"));

        // -- Main container --
        let containers = &pod_spec.containers;
        assert_eq!(containers.len(), 1);
        let main = &containers[0];
        assert_eq!(main.name, "training");
        assert_eq!(
            main.image.as_deref().unwrap(),
            "registry.example.com/foxhunt-training-runtime:v1"
        );

        // Command uses /binaries/ path
        let cmd = main.command.as_ref().unwrap();
        assert_eq!(cmd[0], "/binaries/train_baseline_supervised");

        let args = main.args.as_ref().unwrap();
        assert!(args.contains(&"--symbol=ES.FUT".to_string()));
        assert!(args.contains(&"--data-dir=/data/futures-baseline".to_string()));
        assert!(args.contains(&"--output-dir=/output".to_string()));
        assert!(args.contains(&"--epochs=50".to_string()));

        // GPU resources
        let res = main.resources.as_ref().unwrap();
        let requests = res.requests.as_ref().unwrap();
        assert_eq!(requests.get("nvidia.com/gpu").unwrap().0, "1");
        assert_eq!(requests.get("cpu").unwrap().0, "4");
        assert_eq!(requests.get("memory").unwrap().0, "16Gi");

        // Volume mounts (training-data + output + binaries)
        let mounts = main.volume_mounts.as_ref().unwrap();
        assert_eq!(mounts.len(), 3);
        assert!(mounts.iter().any(|m| m.name == "training-data" && m.mount_path == "/data"));
        assert!(mounts.iter().any(|m| m.name == "output" && m.mount_path == "/output"));
        assert!(mounts.iter().any(|m| m.name == "binaries" && m.mount_path == "/binaries"));

        // -- Volumes (PVC + 2 emptyDir) --
        let volumes = pod_spec.volumes.as_ref().unwrap();
        assert_eq!(volumes.len(), 3);

        let data_vol = volumes.iter().find(|v| v.name == "training-data").unwrap();
        assert_eq!(
            data_vol.persistent_volume_claim.as_ref().unwrap().claim_name,
            "training-data-pvc"
        );

        let output_vol = volumes.iter().find(|v| v.name == "output").unwrap();
        assert!(output_vol.empty_dir.is_some(), "output should be emptyDir");

        let binaries_vol = volumes.iter().find(|v| v.name == "binaries").unwrap();
        assert!(binaries_vol.empty_dir.is_some(), "binaries should be emptyDir");
    }

    #[test]
    fn test_sanitize_k8s_name() {
        assert_eq!(sanitize_k8s_name("training-tft-abc"), "training-tft-abc");
        assert_eq!(
            sanitize_k8s_name("Training-TFT-ABC"),
            "training-tft-abc"
        );
        assert_eq!(
            sanitize_k8s_name("name_with_underscores"),
            "name-with-underscores"
        );

        // Truncation to 63 chars
        let long_name = "a".repeat(100);
        let sanitized = sanitize_k8s_name(&long_name);
        assert!(sanitized.len() <= 63);

        // Leading/trailing hyphens stripped
        assert_eq!(sanitize_k8s_name("--hello--"), "hello");
    }

    #[test]
    fn test_build_job_spec_rl_model() {
        let config = test_config();
        let params = TrainingJobParams {
            job_id: Uuid::parse_str("550e8400-e29b-41d4-a716-446655440000").unwrap(),
            model_type: "dqn".to_string(),
            symbol: "NQ.FUT".to_string(),
            epochs: 100,
            binary: training_binary_for_model("dqn").to_string(),
            extra_args: vec![],
        };

        let job = build_job_spec(&config, &params);
        let name = job.metadata.name.as_deref().unwrap();
        assert!(name.starts_with("training-dqn-"));

        let pod_spec = job.spec.unwrap().template.spec.unwrap();
        let main = &pod_spec.containers[0];
        let cmd = main.command.as_ref().unwrap();
        assert_eq!(cmd[0], "/binaries/train_baseline_rl");

        let args = main.args.as_ref().unwrap();
        assert!(args.contains(&"--symbol=NQ.FUT".to_string()));
        assert!(args.contains(&"--epochs=100".to_string()));
    }

    #[test]
    fn test_hyperopt_binary_for_model() {
        assert_eq!(hyperopt_binary_for_model("dqn"), "hyperopt_baseline_rl");
        assert_eq!(hyperopt_binary_for_model("DQN"), "hyperopt_baseline_rl");
        assert_eq!(hyperopt_binary_for_model("ppo"), "hyperopt_baseline_rl");
        assert_eq!(hyperopt_binary_for_model("PPO"), "hyperopt_baseline_rl");
        assert_eq!(hyperopt_binary_for_model("tft"), "hyperopt_baseline_supervised");
        assert_eq!(hyperopt_binary_for_model("mamba2"), "hyperopt_baseline_supervised");
        assert_eq!(hyperopt_binary_for_model("liquid"), "hyperopt_baseline_supervised");
        assert_eq!(hyperopt_binary_for_model("kan"), "hyperopt_baseline_supervised");
    }

    #[test]
    fn test_build_job_spec_hyperopt() {
        let config = test_config();
        let params = TrainingJobParams {
            job_id: Uuid::parse_str("550e8400-e29b-41d4-a716-446655440000").unwrap(),
            model_type: "dqn".to_string(),
            symbol: "ES.FUT".to_string(),
            epochs: 10,
            binary: hyperopt_binary_for_model("dqn").to_string(),
            extra_args: vec![
                "--output=/output/hyperopt_results.json".to_string(),
                "--model=dqn".to_string(),
                "--trials=20".to_string(),
                "--epochs=10".to_string(),
                "--parallel=0".to_string(),
            ],
        };

        let job = build_job_spec(&config, &params);

        // Job-type label should be hyperopt
        let labels = job.metadata.labels.as_ref().unwrap();
        assert_eq!(labels.get("foxhunt/job-type").unwrap(), "hyperopt");

        let pod_spec = job.spec.unwrap().template.spec.unwrap();
        let main = &pod_spec.containers[0];
        let cmd = main.command.as_ref().unwrap();
        assert_eq!(cmd[0], "/binaries/hyperopt_baseline_rl");

        let args = main.args.as_ref().unwrap();
        assert!(args.contains(&"--symbol=ES.FUT".to_string()));
        assert!(args.contains(&"--trials=20".to_string()));
        assert!(args.contains(&"--model=dqn".to_string()));
        // Should NOT have default training args
        assert!(!args.contains(&"--output-dir=/output".to_string()));
    }

    #[test]
    fn test_from_env_defaults() {
        // Set required env vars, clear optional ones to test defaults.
        std::env::set_var("S3_ENDPOINT", "http://minio:9000");
        std::env::set_var("TRAINING_RUNTIME_IMAGE", "test-runtime:latest");
        std::env::set_var("CALLBACK_ENDPOINT", "http://callback:50053");
        std::env::remove_var("K8S_NAMESPACE");
        std::env::remove_var("S3_BUCKET");
        std::env::remove_var("BINARIES_BUCKET");
        std::env::remove_var("TRAINING_DATA_PVC");
        std::env::remove_var("S3_REGION");

        let cfg = DispatcherConfig::from_env();
        assert_eq!(cfg.namespace, "foxhunt");
        assert_eq!(cfg.runtime_image, "test-runtime:latest");
        assert_eq!(cfg.data_pvc, "training-data-pvc");
        assert_eq!(cfg.binaries_bucket, "foxhunt-binaries");
        assert_eq!(cfg.s3_bucket, "foxhunt-models");
        assert_eq!(cfg.s3_endpoint, "http://minio:9000");
        assert_eq!(cfg.s3_region, "fr-par");
        assert_eq!(cfg.callback_endpoint, "http://callback:50053");
    }

}
