leindex 2.0.3

use super::*;
// Batch/sequence tuning consts live in `runtime_env` (kept out of `runtime.rs`
// for the Large-File gate); import them explicitly here.
use crate::embed::protocol::EmbedRequest;
use crate::embed::runtime_env::{
    DEFAULT_DYNAMIC_ONNX_INFERENCE_BATCH_SIZE, DEFAULT_MIN_AVAILABLE_MB,
    DEFAULT_ONNX_INFERENCE_BATCH_SIZE, MAX_ONNX_SEQUENCE_LEN, ONNX_INFERENCE_BATCH_SIZE_ENV,
    ONNX_SEQUENCE_LEN_ENV,
};
use std::io::Cursor;
use std::sync::Mutex as StdMutex;

static ENV_LOCK: StdMutex<()> = StdMutex::new(());

/// RAII guard that restores an environment variable to its original value on drop.
struct EnvVarGuard {
    key: &'static str,
    original: Option<String>,
}

impl EnvVarGuard {
    fn set(key: &'static str, value: &str) -> Self {
        let original = std::env::var(key).ok();
        unsafe {
            std::env::set_var(key, value);
        }
        Self { key, original }
    }

    fn remove(key: &'static str) -> Self {
        let original = std::env::var(key).ok();
        unsafe {
            std::env::remove_var(key);
        }
        Self { key, original }
    }
}

impl Drop for EnvVarGuard {
    fn drop(&mut self) {
        match &self.original {
            Some(val) => unsafe {
                std::env::set_var(self.key, val);
            },
            None => unsafe {
                std::env::remove_var(self.key);
            },
        }
    }
}

/// Config whose model names resolve to no on-disk file, so
/// `WorkerRuntime::new` skips the ~300s MIGraphX JIT compile AND the lazy
/// reranker cannot load. The worker tests below exercise pooling / idle-timer /
/// dispatch-without-a-session logic, never real inference, so no model is
/// needed. Without this, a real `qwen3-embed-0.6b.onnx` under
/// `~/.leindex/models` makes every `WorkerRuntime::new` compile the model and
/// OOM the test binary (regression introduced when the static model shipped).
/// The rerank model name is likewise poisoned so `ensure_rerank_session` fails
/// at model resolution and dispatch returns an Error frame deterministically —
/// otherwise a host with the rerank model installed would build a real session
/// and the dispatch-error assertions below would be environment-dependent.
fn no_compile_config() -> RuntimeConfig {
    RuntimeConfig {
        model_name: "__leindex_test_no_model__".to_string(),
        rerank_model_name: "__leindex_test_no_rerank_model__".to_string(),
        ..RuntimeConfig::default()
    }
}

#[test]
fn test_runtime_config_default() {
    let config = RuntimeConfig::default();
    assert_eq!(
        config.idle_timeout,
        Duration::from_secs(DEFAULT_IDLE_TIMEOUT_SECS)
    );
    assert_eq!(config.max_frame_size, 16 * 1024 * 1024);
    assert_eq!(config.max_text_size, 1024 * 1024);
    assert_eq!(config.embedding_dim, 1024);
}

#[test]
fn onnx_inference_batch_size_defaults_to_fixed_batch_safe_value() {
    let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
    let _env = EnvVarGuard::remove(ONNX_INFERENCE_BATCH_SIZE_ENV);

    assert_eq!(
        configured_onnx_inference_batch_size("qwen3-embed-0.6b", "cpu"),
        DEFAULT_ONNX_INFERENCE_BATCH_SIZE
    );
    assert_eq!(
        configured_onnx_inference_batch_size("qwen3-embed-0.6b", "cpu"),
        1
    );
    // MIGraphX compiles one fixed shape for every model, including the
    // statically exported b8-s128 qwen3-embed-0.6b graph; a batch-1 policy
    // could never match its compiled .mxr cache.
    assert_eq!(
        configured_onnx_inference_batch_size("qwen3-embed-0.6b", "migraphx"),
        DEFAULT_MIGRAPHX_INFERENCE_BATCH_SIZE
    );
}

#[test]
fn onnx_inference_batch_size_uses_positive_env_override() {
    let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
    let _env = EnvVarGuard::set(ONNX_INFERENCE_BATCH_SIZE_ENV, "32");

    assert_eq!(
        configured_onnx_inference_batch_size("qwen3-embed-0.6b", "migraphx"),
        32
    );
}

#[test]
fn onnx_inference_batch_size_rejects_zero_and_bad_values() {
    let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());

    // Use a single guard capturing the original; intermediate mutations are fine.
    let _env = EnvVarGuard::set(ONNX_INFERENCE_BATCH_SIZE_ENV, "0");
    assert_eq!(
        configured_onnx_inference_batch_size("qwen3-embed-0.6b", "cpu"),
        DEFAULT_ONNX_INFERENCE_BATCH_SIZE
    );

    drop(_env);
    let _env = EnvVarGuard::set(ONNX_INFERENCE_BATCH_SIZE_ENV, "nope");
    assert_eq!(
        configured_onnx_inference_batch_size("qwen3-embed-0.6b", "cpu"),
        DEFAULT_ONNX_INFERENCE_BATCH_SIZE
    );
}

#[cfg(feature = "onnx")]
#[test]
fn qwen_pooling_uses_last_unpadded_token() {
    let runtime = WorkerRuntime::new(no_compile_config());
    let pooled = runtime
        .pool_and_normalize(
            &[
                1.0, 0.0, // first token
                0.0, 2.0, // final real token
                8.0, 8.0, // padding
            ],
            1,
            3,
            &[1, 1, 0],
            2,
        )
        .unwrap();

    assert_eq!(pooled.vectors, vec![0.0, 1.0]);
}

#[cfg(feature = "onnx")]
#[test]
fn qwen_pooling_rejects_short_embedding_output() {
    let runtime = WorkerRuntime::new(no_compile_config());
    let error = runtime
        .pool_and_normalize(&[1.0], 1, 2, &[1, 1], 2)
        .unwrap_err();

    assert_eq!(error.kind, ErrorKind::Inference);
    assert!(error.message.contains("embedding output is too short"));
}

#[test]
fn position_ids_repeat_sequence_for_each_batch_row() {
    assert_eq!(build_position_ids(2, 4), vec![0, 1, 2, 3, 0, 1, 2, 3]);
}

#[test]
fn onnx_sequence_len_defaults_and_clamps_env_override() {
    let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
    let _env = EnvVarGuard::remove(ONNX_SEQUENCE_LEN_ENV);
    assert_eq!(configured_onnx_sequence_len(), DEFAULT_MAX_SEQ_LEN);
    drop(_env);

    let _env = EnvVarGuard::set(ONNX_SEQUENCE_LEN_ENV, "4");
    assert_eq!(configured_onnx_sequence_len(), DEFAULT_MAX_SEQ_LEN);
    drop(_env);

    let _env = EnvVarGuard::set(ONNX_SEQUENCE_LEN_ENV, "256");
    assert_eq!(configured_onnx_sequence_len(), 256);
    drop(_env);

    let _env = EnvVarGuard::set(ONNX_SEQUENCE_LEN_ENV, "4096");
    assert_eq!(configured_onnx_sequence_len(), MAX_ONNX_SEQUENCE_LEN);
}

#[test]
fn dynamic_qwen_uses_batched_inference_by_default() {
    let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
    let _env = EnvVarGuard::remove(ONNX_INFERENCE_BATCH_SIZE_ENV);

    assert_eq!(
        configured_onnx_inference_batch_size("qwen3-embed-0.6b-dynamic", "cpu"),
        DEFAULT_DYNAMIC_ONNX_INFERENCE_BATCH_SIZE
    );
    const _: () = assert!(DEFAULT_DYNAMIC_ONNX_INFERENCE_BATCH_SIZE > 1);
}

#[test]
fn migraphx_uses_one_stable_batch_shape_by_default() {
    let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
    let _env = EnvVarGuard::remove(ONNX_INFERENCE_BATCH_SIZE_ENV);

    assert_eq!(
        configured_onnx_inference_batch_size("qwen3-embed-0.6b-dynamic", "migraphx"),
        DEFAULT_MIGRAPHX_INFERENCE_BATCH_SIZE
    );
}

#[test]
fn test_batch_size_for_dynamic_uint8() {
    let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
    let _env = EnvVarGuard::remove(ONNX_INFERENCE_BATCH_SIZE_ENV);

    // The -dynamic-uint8 suffix must trigger the dynamic batch path, same as
    // -dynamic. MIGraphX gets the stable compiled batch size (8).
    assert_eq!(
        configured_onnx_inference_batch_size("qwen3-embed-0.6b-dynamic-uint8", "migraphx"),
        DEFAULT_MIGRAPHX_INFERENCE_BATCH_SIZE
    );
    // CPU/CUDA gets the larger dynamic batch size (32).
    assert_eq!(
        configured_onnx_inference_batch_size("qwen3-embed-0.6b-dynamic-uint8", "cpu"),
        DEFAULT_DYNAMIC_ONNX_INFERENCE_BATCH_SIZE
    );
}

#[cfg(feature = "onnx")]
#[test]
fn test_extract_u8_dequantization() {
    // Verify the dequantization formula: (value - zero_point) * scale
    // with scale=0.0027450980 and zero_point=109.
    // These are the QuantizeLinear parameters from the electroglyph uint8 model.
    const SCALE: f32 = 0.002_745_098;
    const ZERO_POINT: f32 = 109.0;

    // Test a few representative uint8 values.
    let test_cases: [(u8, f32); 4] = [
        // (input u8, expected dequantized f32)
        (109, 0.0),                     // zero_point -> 0.0
        (0, (0.0 - 109.0) * SCALE),     // min uint8
        (255, (255.0 - 109.0) * SCALE), // max uint8
        (128, (128.0 - 109.0) * SCALE), // mid-range
    ];

    for (input, expected) in test_cases {
        let dequantized = (input as f32 - ZERO_POINT) * SCALE;
        assert!(
            (dequantized - expected).abs() < 1e-6,
            "u8 value {}: expected {}, got {}",
            input,
            expected,
            dequantized
        );
    }
}

#[cfg(feature = "onnx")]
#[test]
fn test_u8_dequant_preserves_unit_norm() {
    // The electroglyph uint8 model L2-normalizes embeddings BEFORE quantizing,
    // then applies QuantizeLinear(scale=0.0027450980, zero_point=109). After
    // dequantization the vector norm should stay close to 1.0 — provided every
    // component is within the quantizer's representable range.
    //
    // With these constants the representable range is
    //   [(0 - 109)*scale, (255 - 109)*scale] = [-0.299, 0.401].
    // So a component like 0.5 is OUT of range and clips to 255 (dequant 0.401),
    // collapsing the norm. A valid unit-norm check must use in-range components.
    // [0.35; 8] has norm ~0.99, all components in range, and stays close to 1.0
    // after quantize+dequant.
    const SCALE: f32 = 0.002_745_098;
    const ZERO_POINT: f32 = 109.0;

    // An 8-dim vector with all components in the quantizer's representable
    // range and near unit norm.
    let original: Vec<f32> = vec![0.35; 8];
    // Quantize: round(value / scale + zero_point), clamp to [0, 255].
    let quantized: Vec<u8> = original
        .iter()
        .map(|&v| {
            let q = (v / SCALE + ZERO_POINT).round() as i32;
            q.clamp(0, 255) as u8
        })
        .collect();
    // Dequantize.
    let dequantized: Vec<f32> = quantized
        .iter()
        .map(|&v| (v as f32 - ZERO_POINT) * SCALE)
        .collect();

    let norm: f32 = dequantized.iter().map(|v| v * v).sum::<f32>().sqrt();
    // The quantization introduces small error, but norm should be near 1.0.
    assert!(
        (norm - 1.0).abs() < 0.1,
        "dequantized vector norm {} should be close to 1.0",
        norm
    );
}

#[test]
fn test_runtime_idle_not_expired_initially() {
    let config = no_compile_config();
    let rt = WorkerRuntime::new(config);
    assert!(!rt.is_idle_expired());
}

#[test]
fn test_runtime_idle_expired_with_zero_timeout() {
    let config = RuntimeConfig {
        idle_timeout: Duration::from_secs(0),
        ..no_compile_config()
    };
    let rt = WorkerRuntime::new(config);
    // With zero timeout, it should be expired immediately
    // (but we need at least a tiny delay for the check)
    std::thread::sleep(Duration::from_millis(1));
    assert!(rt.is_idle_expired());
}

#[test]
fn test_runtime_touch_resets_idle() {
    let config = RuntimeConfig {
        idle_timeout: Duration::from_millis(10),
        ..no_compile_config()
    };
    let rt = WorkerRuntime::new(config);

    std::thread::sleep(Duration::from_millis(20));
    assert!(rt.is_idle_expired());

    rt.touch();
    assert!(!rt.is_idle_expired());
}

#[test]
fn cloned_runtime_shares_idle_activity() {
    fn assert_send_sync<T: Send + Sync>() {}
    assert_send_sync::<WorkerRuntime>();

    let runtime = WorkerRuntime::new(no_compile_config());
    let cloned = runtime.clone();
    assert!(Arc::ptr_eq(&runtime.last_activity, &cloned.last_activity));
    cloned.touch();
    assert!(!runtime.is_idle_expired());
}

#[test]
fn test_shutdown_flag() {
    let config = no_compile_config();
    let rt = WorkerRuntime::new(config);
    let flag = rt.shutdown_flag();

    assert!(!flag.load(Ordering::Relaxed));
    flag.store(true, Ordering::Relaxed);
    assert!(flag.load(Ordering::Relaxed));
}

#[test]
fn test_truncate_text_within_limit() {
    let config = no_compile_config();
    let rt = WorkerRuntime::new(config);
    let text = "hello world".to_string();
    let result = rt.truncate_text(text.clone());
    assert_eq!(result, text);
}

#[test]
fn test_truncate_text_exceeds_limit() {
    let config = RuntimeConfig {
        max_text_size: 10,
        ..no_compile_config()
    };
    let rt = WorkerRuntime::new(config);
    let text = "hello world, this is a long string".to_string();
    let result = rt.truncate_text(text);
    assert!(result.len() <= 10);
    assert_eq!(result, "hello worl");
}

#[test]
fn test_truncate_text_unicode_boundary() {
    let config = RuntimeConfig {
        max_text_size: 10,
        ..no_compile_config()
    };
    let rt = WorkerRuntime::new(config);
    // "héllo" has multi-byte chars
    let text = "héllo wörld test".to_string();
    let result = rt.truncate_text(text);
    assert!(result.len() <= 10);
    // Should not panic and should be valid UTF-8
    assert!(result.is_char_boundary(result.len()));
}

#[test]
fn test_handle_embed_empty_batch() {
    let config = no_compile_config();
    let rt = WorkerRuntime::new(config);

    let request = EmbedRequest {
        texts: vec![],
        expected_dim: 1024,
        cache_keys: vec![],
    };
    let frame = protocol::embed_request_frame(BatchId::new(1), request).unwrap();
    let result = rt.handle_embed(&frame, &Arc::new(AtomicBool::new(false)));

    // Empty batch returns Ok early (before any ONNX session check),
    // so .unwrap() is safe regardless of feature flag.
    let response = result.unwrap();
    assert_eq!(response.count, 0);
    assert_eq!(response.dimension, 1024);
    assert!(response.vectors.is_empty());
}

#[test]
fn test_handle_embed_returns_flat_row_major() {
    let config = no_compile_config();
    let rt = WorkerRuntime::new(config);

    let request = EmbedRequest {
        texts: vec!["hello".to_string(), "world".to_string()],
        expected_dim: 8,
        cache_keys: vec![],
    };
    let frame = protocol::embed_request_frame(BatchId::new(1), request).unwrap();
    let result = rt.handle_embed(&frame, &Arc::new(AtomicBool::new(false)));

    // When ORT or the model is unavailable (no model on disk, ORT not
    // discovered, etc.), the worker returns ModelNotFound. On a developer
    // machine that has both /usr/local/lib/libonnxruntime.so and a real
    // model in `~/.leindex/models/`, ORT inference may actually run and
    // fail with a different error (Inference); treat that as acceptable
    // since the contract under test is "no crash, structured error".
    #[cfg(feature = "onnx")]
    {
        let err = result.unwrap_err();
        assert!(
            err.kind == ErrorKind::ModelNotFound || err.kind == ErrorKind::Inference,
            "expected ModelNotFound or Inference, got {:?}: {}",
            err.kind,
            err.message
        );
    }

    // Without ONNX feature, returns zero vectors
    #[cfg(not(feature = "onnx"))]
    {
        let response = result.unwrap();
        assert_eq!(response.count, 2);
        assert_eq!(response.dimension, 8);
        assert_eq!(response.vectors.len(), 16);
        assert_eq!(response.get_embedding(0).unwrap().len(), 8);
        assert_eq!(response.get_embedding(1).unwrap().len(), 8);
    }
}

#[test]
fn test_handle_embed_preserves_ordering() {
    let config = no_compile_config();
    let rt = WorkerRuntime::new(config);

    let texts: Vec<String> = (0..5).map(|i| format!("text {}", i)).collect();
    let request = EmbedRequest {
        texts: texts.clone(),
        expected_dim: 4,
        cache_keys: vec![],
    };
    let frame = protocol::embed_request_frame(BatchId::new(1), request).unwrap();
    let result = rt.handle_embed(&frame, &Arc::new(AtomicBool::new(false)));

    // Same rationale as test_handle_embed_returns_flat_row_major: developer
    // machines with ORT + a real model present may reach inference and
    // surface an Inference error instead of ModelNotFound. Both are
    // acceptable; the contract under test is "no crash, structured error".
    #[cfg(feature = "onnx")]
    {
        let err = result.unwrap_err();
        assert!(
            err.kind == ErrorKind::ModelNotFound || err.kind == ErrorKind::Inference,
            "expected ModelNotFound or Inference, got {:?}: {}",
            err.kind,
            err.message
        );
    }

    // Without ONNX feature, returns zero vectors with correct count
    #[cfg(not(feature = "onnx"))]
    {
        let response = result.unwrap();
        assert_eq!(response.count, 5);
        for i in 0..5 {
            assert!(response.get_embedding(i).is_some());
        }
    }
}

#[test]
fn test_dispatch_embed_request() {
    let config = no_compile_config();
    let rt = WorkerRuntime::new(config);

    let request = EmbedRequest {
        texts: vec!["test".to_string()],
        expected_dim: 4,
        cache_keys: vec![],
    };
    let frame = protocol::embed_request_frame(BatchId::new(42), request).unwrap();
    let response_frame = rt.dispatch(&frame);

    assert_eq!(response_frame.header.batch_id, BatchId::new(42));

    // Without a real ONNX session, dispatch returns an error frame
    #[cfg(feature = "onnx")]
    {
        assert_eq!(response_frame.header.msg_type, MsgType::Error);
    }

    // Without ONNX feature, dispatch returns a success response
    #[cfg(not(feature = "onnx"))]
    {
        assert_eq!(response_frame.header.msg_type, MsgType::EmbedResponse);
    }
}

#[test]
fn test_dispatch_rerank_request() {
    let config = no_compile_config();
    let rt = WorkerRuntime::new(config);

    let request = protocol::RerankRequest {
        query: "test".to_string(),
        documents: vec![protocol::RerankDocument {
            id: "doc1".to_string(),
            content: "content".to_string(),
            initial_score: 0.9,
        }],
    };
    let frame = protocol::rerank_request_frame(BatchId::new(7), request).unwrap();
    let response_frame = rt.dispatch(&frame);

    assert_eq!(response_frame.header.batch_id, BatchId::new(7));

    // Without a real ONNX session, dispatch returns an error frame
    #[cfg(feature = "onnx")]
    {
        assert_eq!(response_frame.header.msg_type, MsgType::Error);
    }

    // Without ONNX feature, dispatch returns a success response
    #[cfg(not(feature = "onnx"))]
    {
        assert_eq!(response_frame.header.msg_type, MsgType::RerankResponse);
    }
}

#[test]
fn test_dispatch_unknown_message_type() {
    let config = no_compile_config();
    let rt = WorkerRuntime::new(config);

    let frame = Frame {
        header: protocol::FrameHeader {
            batch_id: BatchId::new(99),
            msg_type: MsgType::Error, // Unexpected from main daemon
        },
        payload: vec![],
    };
    let response_frame = rt.dispatch(&frame);

    assert_eq!(response_frame.header.batch_id, BatchId::new(99));
    assert_eq!(response_frame.header.msg_type, MsgType::Error);
}

#[test]
fn test_run_loop_single_request() {
    let config = RuntimeConfig {
        idle_timeout: Duration::from_secs(300),
        ..no_compile_config()
    };
    let rt = WorkerRuntime::new(config);

    // Build a single embed request frame
    let request = EmbedRequest {
        texts: vec!["hello".to_string()],
        expected_dim: 4,
        cache_keys: vec![],
    };
    let frame = protocol::embed_request_frame(BatchId::new(1), request).unwrap();
    let wire = frame.encode_wire().unwrap();

    // Create a reader that will return the frame then EOF
    let reader = Cursor::new(wire);
    let writer = Cursor::new(Vec::<u8>::new());

    let result = rt.run_loop(reader, writer);
    assert!(result.is_ok());
}

// ── Task 7: provider selection precedence & truthfulness ──────────────

#[test]
fn runtime_config_provider_env_overrides_toml_and_default() {
    // Direct-worker precedence: env > TOML > "auto". When the env var is set,
    // it wins regardless of what the TOML contains (the TOML is read via the
    // process-global OnceLock and may hold any value on the test host).
    let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
    let _env = EnvVarGuard::set("LEINDEX_WORKER_EXECUTION_PROVIDER", "cuda");

    let config = RuntimeConfig::from_env();
    assert_eq!(
        config.execution_provider, "cuda",
        "env var must take precedence over TOML and the 'auto' default"
    );

    drop(_env);
    // With the env var unset, from_env falls back to TOML or "auto". Either
    // way it must be a non-empty lowercase value (normalized).
    let _env_removed = EnvVarGuard::remove("LEINDEX_WORKER_EXECUTION_PROVIDER");
    let config = RuntimeConfig::from_env();
    assert!(
        !config.execution_provider.trim().is_empty(),
        "provider must always resolve to a non-empty value"
    );
    // The value handed to the selector is already normalized lowercased.
    assert_eq!(
        config.execution_provider,
        config.execution_provider.to_ascii_lowercase(),
    );
}

#[test]
fn runtime_config_provider_blanks_env_falls_through() {
    // A blank/whitespace env value is treated as unset → falls through to TOML
    // or "auto". This prevents a stray `LEINDEX_WORKER_EXECUTION_PROVIDER=""`
    // from producing an empty provider string that bypasses normalization.
    let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
    let _env = EnvVarGuard::set("LEINDEX_WORKER_EXECUTION_PROVIDER", "   ");
    let config = RuntimeConfig::from_env();
    assert!(
        !config.execution_provider.trim().is_empty(),
        "blank env value must fall through to a non-empty default"
    );
}

#[test]
fn runtime_config_min_available_mb_defaults_to_documented_floor() {
    // Codex P1: with LEINDEX_WORKER_MIN_AVAILABLE_MB unset, from_env applies
    // the documented 2048 MiB default so the refusal guard is active in the
    // default configuration — an unset variable must not silently bypass it.
    let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
    let _env = EnvVarGuard::remove("LEINDEX_WORKER_MIN_AVAILABLE_MB");
    let config = RuntimeConfig::from_env();
    assert_eq!(
        config.min_available_mb,
        Some(DEFAULT_MIN_AVAILABLE_MB),
        "unset MIN_AVAILABLE_MB must apply the documented 2048 MiB default"
    );

    // An explicit positive value wins.
    let _env2 = EnvVarGuard::set("LEINDEX_WORKER_MIN_AVAILABLE_MB", "4096");
    let config = RuntimeConfig::from_env();
    assert_eq!(config.min_available_mb, Some(4096));

    // An explicit 0 is treated as disabled (None) per the env contract.
    let _env3 = EnvVarGuard::set("LEINDEX_WORKER_MIN_AVAILABLE_MB", "0");
    let config = RuntimeConfig::from_env();
    assert_eq!(config.min_available_mb, None, "0 = disabled");

    // Codex P2: a malformed override must NOT silently bypass the guard —
    // fall back to the documented floor rather than resolving to None like
    // an explicit 0.
    let _env4 = EnvVarGuard::set("LEINDEX_WORKER_MIN_AVAILABLE_MB", "not-a-number");
    let config = RuntimeConfig::from_env();
    assert_eq!(
        config.min_available_mb,
        Some(DEFAULT_MIN_AVAILABLE_MB),
        "malformed MIN_AVAILABLE_MB must keep the documented 2048 MiB default"
    );
}

#[test]
fn explicit_gpu_provider_select_returns_cpu_fallback_not_hard_error() {
    // Explicit-GPU→CPU is the documented neural-fallback path: select()
    // returns Err(ProviderSelection) whose provider is CPU, so the runtime
    // can log a warning and continue. This is the contract that
    // try_provider_or_cpu preserves when error_on_failure() turns a failed
    // registration into an Err — the fallback rebuild must NOT be a hard error.
    use crate::embed::provider::ExecutionProviderSelector;
    let result = ExecutionProviderSelector::select("cuda");
    match result {
        Ok(selection) => {
            // CUDA is available on this host — still a concrete provider, never
            // "auto", and registration will be attempted with error_on_failure.
            let name = selection.name();
            assert!(
                name == "cuda" || name == "cpu",
                "unexpected cuda-resolution: {name}"
            );
        }
        Err(fallback) => {
            assert_eq!(fallback.fallback_name(), "cpu");
            assert!(!fallback.is_requested_provider());
            assert!(
                fallback.reason().contains("CUDA"),
                "fallback reason should name the missing provider: {}",
                fallback.reason()
            );
        }
    }
}

#[cfg(feature = "onnx")]
#[test]
fn runtime_resolved_provider_is_never_auto_for_embed_or_rerank() {
    // End-to-end invariant: after WorkerRuntime construction, the provider
    // recorded on provider_runtime_status (used by both the embedder session
    // and the lazy reranker via ensure_rerank_session) is never the unresolved
    // "auto" token. It is always one of cpu/cuda/migraphx/coreml.
    //
    // no_compile_config() points at a non-existent model so build_session
    // fails at commit_from_file; the provider_runtime_status then retains the
    // pre-build "cpu" fallback value, which is itself a valid concrete token.
    // The invariant under test is that NO path surfaces "auto" to a session
    // builder or the health response.
    let rt = WorkerRuntime::new(RuntimeConfig {
        execution_provider: "auto".to_string(),
        ..no_compile_config()
    });
    let health = rt.health_response(crate::embed::protocol::WorkerState::Initializing, None);
    let provider = health.provider.expect("health response carries a provider");
    assert_ne!(
        provider, "auto",
        "the provider handed to embedder/reranker session builders must be concrete"
    );
    assert!(
        matches!(provider.as_str(), "cpu" | "cuda" | "migraphx" | "coreml"),
        "resolved provider must be one of the concrete tokens, got {provider}"
    );
}

#[test]
fn test_run_loop_multiple_requests_same_runtime() {
    // VAL-CPHASE-006: Worker remains reusable across successive batches
    let config = RuntimeConfig {
        idle_timeout: Duration::from_secs(300),
        ..no_compile_config()
    };
    let rt = WorkerRuntime::new(config);

    // Build two embed request frames
    let request1 = EmbedRequest {
        texts: vec!["first".to_string()],
        expected_dim: 4,
        cache_keys: vec![],
    };
    let request2 = EmbedRequest {
        texts: vec!["second".to_string()],
        expected_dim: 4,
        cache_keys: vec![],
    };

    let frame1 = protocol::embed_request_frame(BatchId::new(1), request1).unwrap();
    let frame2 = protocol::embed_request_frame(BatchId::new(2), request2).unwrap();

    let wire1 = frame1.encode_wire().unwrap();
    let wire2 = frame2.encode_wire().unwrap();

    let mut combined = wire1.clone();
    combined.extend_from_slice(&wire2);

    let reader = Cursor::new(combined);
    let writer = Cursor::new(Vec::<u8>::new());

    let result = rt.run_loop(reader, writer);
    assert!(result.is_ok());

    // Verify both responses were written
    result.unwrap();
}

#[test]
fn test_idle_timeout_causes_exit() {
    // VAL-CPHASE-007: Worker tears down on idle
    let config = RuntimeConfig {
        idle_timeout: Duration::from_millis(1),
        ..no_compile_config()
    };
    let rt = WorkerRuntime::new(config);

    // Empty input — the loop should detect idle timeout
    let reader = Cursor::new(Vec::<u8>::new());
    let writer = Cursor::new(Vec::<u8>::new());

    // This will fail because there's no data to read, but the idle check
    // happens before the read. However, with empty input, read_exact will
    // return UnexpectedEof immediately, which is a clean shutdown.
    let result = rt.run_loop(reader, writer);
    assert!(result.is_ok());
}

// ── VAL-ONNX embed batch loop ─────────────────────────────────────────
// These tests exercise `run_onnx_embed_batch_loop` directly with a mocked
// sub-batch runner (injecting a deterministic pooled vector per row) so the
// provider batching/trimming logic is verified without needing a real ONNX
// model. They cover the CRITICAL missing-else-branch bug (VAL-ONNX-001),
// fixed-batch full sub-batches (VAL-ONNX-002), padding+trim for fixed-batch
// partial sub-batches (VAL-ONNX-003), and the EmbedResponse count/dimension
// invariant across providers and counts (VAL-ONNX-005, VAL-ONNX-006).

#[cfg(feature = "onnx")]
fn test_encoding(marker: u64) -> tokenizers::Encoding {
    use std::collections::HashMap;
    tokenizers::Encoding::new(
        vec![marker as u32, (marker + 1) as u32],
        vec![0, 0],
        vec![marker.to_string(), (marker + 1).to_string()],
        vec![None, None],
        vec![(0, 1), (1, 2)],
        vec![0, 0],
        vec![1, 1],
        vec![],
        HashMap::new(),
    )
}

#[cfg(feature = "onnx")]
fn embed_encodings(n: usize) -> Vec<tokenizers::Encoding> {
    (0..n).map(|i| test_encoding(i as u64)).collect()
}

#[cfg(feature = "onnx")]
#[test]
fn test_embed_batch_cpu_all_sub_batches_processed() {
    // VAL-ONNX-001: CPU/CUDA (fixed_batch == false) must process EVERY
    // sub-batch, not silently skip them (the historical missing-else-branch
    // bug). For N > inference_batch_size the loop should call the runner for
    // each chunk and concatenate all rows.
    let rt = WorkerRuntime::new(no_compile_config());
    let encodings = embed_encodings(17);
    let batch_size = 8usize;
    let dim = 4usize;

    let mut batches: Vec<usize> = Vec::new();
    let all_pooled = rt
        .run_onnx_embed_batch_loop(&encodings, batch_size, false, dim, |sub, dim| {
            batches.push(sub.len());
            let mut out = Vec::with_capacity(sub.len() * dim);
            for (i, _) in sub.iter().enumerate() {
                out.extend(std::iter::repeat_n(i as f32, dim));
            }
            Ok(out)
        })
        .unwrap();

    // 17 encodings chunked by 8 => [8, 8, 1]; all sub-batches ran.
    assert_eq!(batches, vec![8, 8, 1]);
    assert_eq!(all_pooled.len(), 17 * dim);
    assert_ne!(all_pooled.len(), 0, "CPU/CUDA embed must not be empty");
}

#[cfg(feature = "onnx")]
#[test]
fn test_embed_batch_migraphx_full_sub_batches_no_padding() {
    // VAL-ONNX-002: fixed-batch provider with an exact multiple of the batch
    // size runs each full sub-batch unchanged — no padding is applied.
    let rt = WorkerRuntime::new(no_compile_config());
    let batch_size = 8usize;
    let dim = 4usize;
    let encodings = embed_encodings(16); // 2 full sub-batches of 8

    let mut batches: Vec<usize> = Vec::new();
    let all_pooled = rt
        .run_onnx_embed_batch_loop(&encodings, batch_size, true, dim, |sub, dim| {
            batches.push(sub.len());
            Ok(vec![1.0f32; sub.len() * dim])
        })
        .unwrap();

    assert_eq!(batches, vec![8, 8]);
    assert_eq!(batches.len(), 2);
    assert_eq!(all_pooled.len(), 16 * dim);
}

#[cfg(feature = "onnx")]
#[test]
fn test_embed_batch_migraphx_partial_sub_batch_padded_and_trimmed() {
    // VAL-ONNX-003: fixed-batch provider with a non-multiple input must pad
    // the final partial sub-batch up to inference_batch_size before running,
    // then TRIM the results back to the real row count. The mocked runner
    // emits row-index-tagged rows, so we can prove the padding rows were cut.
    let rt = WorkerRuntime::new(no_compile_config());
    let batch_size = 8usize;
    let dim = 4usize;
    let encodings = embed_encodings(10); // 8 + 2 partial

    let mut batches: Vec<usize> = Vec::new();
    let all_pooled = rt
        .run_onnx_embed_batch_loop(&encodings, batch_size, true, dim, |sub, dim| {
            batches.push(sub.len());
            // Rows tagged with their 0-based index within the passed sub-batch.
            let mut out = Vec::with_capacity(sub.len() * dim);
            for (i, _) in sub.iter().enumerate() {
                out.extend(std::iter::repeat_n(i as f32, dim));
            }
            Ok(out)
        })
        .unwrap();

    // The runner is called with an 8-sized padded batch for the 2-row tail.
    assert_eq!(batches, vec![8, 8]);
    // 10 real rows remain after trimming the padded (2-row) sub-batch.
    assert_eq!(all_pooled.len(), 10 * dim);
    // The trimmed tail rows must be the first 2 rows (indices 0 and 1) of the
    // padded output, whose tags are 0.0 and 1.0 — not the padding rows 2..7.
    let tail = &all_pooled[all_pooled.len() - dim..];
    assert_eq!(
        tail, &[1.0f32; 4],
        "last row must be the real row 1, not padding"
    );
    // And the very last element equals real-row tag 1.0 (padding would be 7.0).
    assert_eq!(*all_pooled.last().unwrap(), 1.0);
}

#[cfg(feature = "onnx")]
#[test]
fn test_embed_response_invariant_all_providers() {
    // VAL-ONNX-005/006: For every (provider, input_count) combination the
    // batch loop must produce flattened vectors of length count*dimension,
    // which EmbedResponse::new (debug_assert_eq!) accepts without panicking.
    let rt = WorkerRuntime::new(no_compile_config());
    let batch_size = 8usize;
    let dim = 4usize;
    let providers: &[(&str, bool)] = &[("cpu", false), ("migraphx", true)];
    let counts: &[usize] = &[1, 3, 8, 9, 16, 17];

    for (provider, fixed_batch) in providers {
        for &count in counts {
            let encodings = embed_encodings(count);
            let mut batches: Vec<usize> = Vec::new();
            let all_pooled = rt
                .run_onnx_embed_batch_loop(&encodings, batch_size, *fixed_batch, dim, |sub, dim| {
                    batches.push(sub.len());
                    Ok(vec![0.5f32; sub.len() * dim])
                })
                .unwrap();
            assert_eq!(
                all_pooled.len(),
                count * dim,
                "{provider} count={count}: vectors.len() != count*dimension"
            );
            // Constructing EmbedResponse::new runs its debug_assert_eq!; if
            // vectors.len() != count*dim it panics, flagging the regression.
            let response = EmbedResponse::new(all_pooled.clone(), count, dim);
            assert_eq!(response.vectors.len(), count * dim);
            assert_eq!(response.count, count);
            assert_eq!(response.dimension, dim);
            assert_eq!(
                response.vectors.len(),
                response.count * response.dimension,
                "EmbedResponse invariant broken for {provider} count={count}"
            );
        }
    }
}

#[cfg(feature = "onnx")]
#[test]
fn test_embed_text_batch_loop_tokenizes_and_infers_per_sub_batch() {
    // Fix B: tokenization must be bounded to the inference batch and inference
    // must begin before later text batches are tokenized.
    let rt = WorkerRuntime::new(no_compile_config());
    let texts: Vec<String> = (0..17).map(|i| format!("text-{i}")).collect();
    let mut tokenized_sizes = Vec::new();
    let events = std::cell::RefCell::new(Vec::new());
    let dim = 4usize;

    let pooled = rt
        .run_onnx_embed_text_batch_loop(
            &texts,
            8,
            false,
            dim,
            &Arc::new(AtomicBool::new(false)),
            |sub_texts| {
                tokenized_sizes.push(sub_texts.len());
                events
                    .borrow_mut()
                    .push(format!("tokenize-{}", sub_texts.len()));
                Ok(embed_encodings(sub_texts.len()))
            },
            |encodings, dim| {
                events
                    .borrow_mut()
                    .push(format!("infer-{}", encodings.len()));
                Ok(vec![1.0f32; encodings.len() * dim])
            },
        )
        .unwrap();

    assert_eq!(tokenized_sizes, vec![8, 8, 1]);
    assert_eq!(pooled.len(), texts.len() * dim);
    assert_eq!(
        events.into_inner(),
        vec![
            "tokenize-8".to_string(),
            "infer-8".to_string(),
            "tokenize-8".to_string(),
            "infer-8".to_string(),
            "tokenize-1".to_string(),
            "infer-1".to_string(),
        ]
    );
}

#[cfg(feature = "onnx")]
#[test]
fn test_embed_text_batch_loop_fixed_batch_pads_after_per_batch_tokenization() {
    let rt = WorkerRuntime::new(no_compile_config());
    let texts: Vec<String> = (0..10).map(|i| format!("text-{i}")).collect();
    let mut tokenized_sizes = Vec::new();
    let mut inferred_sizes = Vec::new();
    let dim = 2usize;

    let pooled = rt
        .run_onnx_embed_text_batch_loop(
            &texts,
            8,
            true,
            dim,
            &Arc::new(AtomicBool::new(false)),
            |sub_texts| {
                tokenized_sizes.push(sub_texts.len());
                Ok(embed_encodings(sub_texts.len()))
            },
            |encodings, dim| {
                inferred_sizes.push(encodings.len());
                Ok(vec![0.5f32; encodings.len() * dim])
            },
        )
        .unwrap();

    assert_eq!(tokenized_sizes, vec![8, 2]);
    assert_eq!(inferred_sizes, vec![8, 8]);
    assert_eq!(pooled.len(), texts.len() * dim);
}

#[cfg(feature = "onnx")]
#[test]
fn test_embed_text_batch_loop_later_tokenizer_error_stops_before_next_inference() {
    let rt = WorkerRuntime::new(no_compile_config());
    let texts: Vec<String> = (0..10).map(|i| format!("text-{i}")).collect();
    let mut inferred_batches = Vec::new();
    let mut tokenizer_calls = 0usize;

    let error = rt
        .run_onnx_embed_text_batch_loop(
            &texts,
            8,
            false,
            2,
            &Arc::new(AtomicBool::new(false)),
            |_sub_texts| {
                tokenizer_calls += 1;
                if tokenizer_calls == 2 {
                    Err(WorkerError {
                        kind: ErrorKind::Tokenizer,
                        message: "synthetic tokenizer failure".to_string(),
                    })
                } else {
                    Ok(embed_encodings(8))
                }
            },
            |encodings, dim| {
                inferred_batches.push(encodings.len());
                Ok(vec![1.0f32; encodings.len() * dim])
            },
        )
        .unwrap_err();

    assert_eq!(inferred_batches, vec![8]);
    assert_eq!(tokenizer_calls, 2);
    assert_eq!(error.kind, ErrorKind::Tokenizer);
    assert!(error.message.contains("synthetic tokenizer failure"));
}

#[cfg(feature = "onnx")]
#[test]
fn test_embed_text_batch_loop_checks_cancel_between_sub_batches() {
    let rt = WorkerRuntime::new(no_compile_config());
    let texts: Vec<String> = (0..10).map(|i| format!("text-{i}")).collect();
    let mut inferred_batches = Vec::new();
    let cancel_token = Arc::new(AtomicBool::new(false));
    let cancel_from_runner = Arc::clone(&cancel_token);

    let error = rt
        .run_onnx_embed_text_batch_loop(
            &texts,
            8,
            false,
            2,
            &cancel_token,
            |_sub_texts| Ok(embed_encodings(8)),
            |encodings, dim| {
                inferred_batches.push(encodings.len());
                cancel_from_runner.store(true, Ordering::Release);
                Ok(vec![1.0f32; encodings.len() * dim])
            },
        )
        .unwrap_err();

    assert_eq!(inferred_batches, vec![8]);
    assert_eq!(error.kind, ErrorKind::Inference);
    assert!(error.message.contains("cancelled"));
}

// ── Batch size provider-precedence fix ────────────────────────────────
// MIGraphX compiles ONE fixed input shape per model, so every model —
// dynamic or statically exported — runs at the fixed MIGraphX batch on
// that provider (the batch loop pads the final partial batch). The
// statically exported qwen3-embed-0.6b graph is itself b8-s128 (verified
// empirically: ORT rejects batch-1 inputs with "Got: 1, Expected: 8"), and
// its compiled .mxr cache is keyed to that shape, so the previous
// batch-1 policy could never use the GPU. CPU keeps the model-specific
// defaults (batch 1 for static exports, 32 for -dynamic variants).

#[test]
fn test_non_dynamic_model_migraphx_returns_fixed_batch() {
    let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
    let _env = EnvVarGuard::remove(ONNX_INFERENCE_BATCH_SIZE_ENV);

    assert_eq!(
        configured_onnx_inference_batch_size("qwen3-embed-0.6b", "migraphx"),
        DEFAULT_MIGRAPHX_INFERENCE_BATCH_SIZE,
        "non-dynamic model with migraphx must return the fixed MIGraphX batch (the static graph is b8-s128)"
    );
    assert_eq!(
        configured_onnx_inference_batch_size("qwen3-embed-0.6b", "migraphx"),
        8
    );
}

#[test]
fn test_non_dynamic_model_rocm_returns_fixed_batch() {
    let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
    let _env = EnvVarGuard::remove(ONNX_INFERENCE_BATCH_SIZE_ENV);

    assert_eq!(
        configured_onnx_inference_batch_size("qwen3-embed-0.6b", "rocm"),
        DEFAULT_MIGRAPHX_INFERENCE_BATCH_SIZE,
        "non-dynamic model with rocm must return the fixed MIGraphX batch"
    );
}

#[test]
fn test_dynamic_model_migraphx_returns_batch_8() {
    let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
    let _env = EnvVarGuard::remove(ONNX_INFERENCE_BATCH_SIZE_ENV);

    assert_eq!(
        configured_onnx_inference_batch_size("qwen3-embed-0.6b-dynamic", "migraphx"),
        DEFAULT_MIGRAPHX_INFERENCE_BATCH_SIZE,
        "dynamic model with migraphx must return the MIGraphX stable batch size"
    );
    assert_eq!(
        configured_onnx_inference_batch_size("qwen3-embed-0.6b-dynamic", "migraphx"),
        8
    );
}

#[test]
fn test_dynamic_model_rocm_returns_batch_8() {
    let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
    let _env = EnvVarGuard::remove(ONNX_INFERENCE_BATCH_SIZE_ENV);

    assert_eq!(
        configured_onnx_inference_batch_size("qwen3-embed-0.6b-dynamic", "rocm"),
        DEFAULT_MIGRAPHX_INFERENCE_BATCH_SIZE,
        "dynamic model with rocm must return the MIGraphX stable batch size"
    );
}

#[test]
fn test_non_dynamic_model_cpu_returns_batch_1() {
    let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
    let _env = EnvVarGuard::remove(ONNX_INFERENCE_BATCH_SIZE_ENV);

    assert_eq!(
        configured_onnx_inference_batch_size("qwen3-embed-0.6b", "cpu"),
        DEFAULT_ONNX_INFERENCE_BATCH_SIZE,
        "non-dynamic model with cpu must return batch_size=1 (unchanged)"
    );
}

#[test]
fn test_dynamic_model_cpu_returns_default_dynamic() {
    let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
    let _env = EnvVarGuard::remove(ONNX_INFERENCE_BATCH_SIZE_ENV);

    assert_eq!(
        configured_onnx_inference_batch_size("qwen3-embed-0.6b-dynamic", "cpu"),
        DEFAULT_DYNAMIC_ONNX_INFERENCE_BATCH_SIZE,
        "dynamic model with cpu must return the dynamic batch size"
    );
}

#[test]
fn test_non_dynamic_model_migraphx_case_insensitive() {
    let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
    let _env = EnvVarGuard::remove(ONNX_INFERENCE_BATCH_SIZE_ENV);

    // Provider matching is case-insensitive.
    assert_eq!(
        configured_onnx_inference_batch_size("qwen3-embed-0.6b", "MIGRAPHX"),
        DEFAULT_MIGRAPHX_INFERENCE_BATCH_SIZE,
        "non-dynamic model with uppercase MIGRAPHX must return the fixed MIGraphX batch"
    );
    assert_eq!(
        configured_onnx_inference_batch_size("qwen3-embed-0.6b", "ROCm"),
        DEFAULT_MIGRAPHX_INFERENCE_BATCH_SIZE,
    );
}

// ── Collapsed batch dimension handling ────────────────────────────────
// When a non-dynamic ONNX model receives batch_size > 1, it may silently
// collapse the batch dimension and return [1, seq_len, hidden_dim] instead
// of [batch_size, seq_len, hidden_dim]. The runtime must detect this and
// retry each sequence individually rather than erroring and triggering
// TF-IDF fallback.

#[cfg(feature = "onnx")]
#[test]
fn test_collapsed_batch_sentinel_is_detectable() {
    // Verify the sentinel constant exists and has the expected prefix.
    assert!(COLLAPSED_BATCH_SENTINEL.starts_with("__"));
    assert!(!COLLAPSED_BATCH_SENTINEL.is_empty());
}

#[cfg(feature = "onnx")]
#[test]
fn test_finalize_embed_output_detects_collapsed_batch() {
    // When the model returns [1, seq_len, hidden_dim] but batch_size > 1 was
    // sent, finalize_embed_output must return an error whose message starts
    // with the COLLAPSED_BATCH_SENTINEL so the caller can retry individually.
    //
    // We cannot easily construct a real SessionOutputs without a model, but
    // we can verify the sentinel-based detection logic by checking that the
    // sentinel prefix is what the retry path matches on.
    let fake_error_msg = format!(
        "{}: model collapsed batch dimension (sent 3, got [1, 128, 1024])",
        COLLAPSED_BATCH_SENTINEL
    );
    assert!(
        fake_error_msg.starts_with(COLLAPSED_BATCH_SENTINEL),
        "collapsed batch error must start with the sentinel for retry detection"
    );
}

#[cfg(feature = "onnx")]
#[test]
fn test_non_collapsed_error_does_not_match_sentinel() {
    // A regular inference error must NOT match the sentinel, so it is not
    // mistaken for a collapsed-batch retry signal.
    let regular_error = "ONNX inference failed: shape mismatch";
    assert!(
        !regular_error.starts_with(COLLAPSED_BATCH_SENTINEL),
        "regular errors must not match the collapsed batch sentinel"
    );
}