inference/

engine.rs

//! Core embedding engine for generating vector embeddings from text.
//!
//! The `EmbeddingEngine` provides a high-level interface for:
//! - Loading ONNX INT8 embedding models from HuggingFace Hub
//! - Generating embeddings for single texts or batches
//! - Automatic batching and parallel processing via ONNX Runtime
//!
//! # Example
//!
//! ```no_run
//! use inference::{EmbeddingEngine, ModelConfig, EmbeddingModel};
//!
//! #[tokio::main]
//! async fn main() {
//!     let config = ModelConfig::new(EmbeddingModel::MiniLM);
//!     let engine = EmbeddingEngine::new(config).await.unwrap();
//!
//!     // Embed a single query
//!     let embedding = engine.embed_query("What is machine learning?").await.unwrap();
//!     println!("Embedding dimension: {}", embedding.len());
//!
//!     // Embed multiple documents
//!     let docs = vec![
//!         "Machine learning is a subset of AI.".to_string(),
//!         "Deep learning uses neural networks.".to_string(),
//!     ];
//!     let embeddings = engine.embed_documents(&docs).await.unwrap();
//!     println!("Generated {} embeddings", embeddings.len());
//! }
//! ```

use crate::batch::{mean_pooling, normalize_embeddings, BatchProcessor};
use crate::error::{InferenceError, Result};
use crate::models::{EmbeddingModel, ModelConfig};
use ort::inputs;
use ort::session::builder::GraphOptimizationLevel;
use ort::session::Session;
use ort::value::Tensor;
use parking_lot::Mutex;
use std::io::Read;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use tokenizers::Tokenizer;
use tracing::{debug, info, instrument, warn};

use ort::execution_providers::CUDAExecutionProvider;

/// The main embedding engine for generating vector embeddings.
///
/// This struct is thread-safe and can be shared across async tasks.
/// ORT sessions are mutex-guarded (run() takes &mut self) and held in a
/// pool so concurrent callers can embed without head-of-line blocking.
/// CPU-heavy inference is offloaded via `tokio::task::spawn_blocking`.
pub struct EmbeddingEngine {
    /// Pool of ONNX Runtime sessions — each guarded independently.
    /// Concurrent callers round-robin across sessions via `next_session`,
    /// eliminating Mutex head-of-line blocking for batch embedding workloads.
    sessions: Vec<Arc<Mutex<Session>>>,
    /// Round-robin counter: atomically incremented per batch, wraps via modulo.
    next_session: AtomicUsize,
    /// Batch processor for tokenization (Arc-wrapped for spawn_blocking)
    processor: Arc<BatchProcessor>,
    /// Model configuration
    config: ModelConfig,
    /// Embedding dimension
    dimension: usize,
}

impl EmbeddingEngine {
    /// Create a new embedding engine with the given configuration.
    ///
    /// Downloads the FP32 ONNX model (`model.onnx`) when CUDA EP is enabled, or the
    /// INT8 quantized model (`model_quantized.onnx`) otherwise.
    #[instrument(skip_all, fields(model = %config.model))]
    pub async fn new(config: ModelConfig) -> Result<Self> {
        // Resolve GPU mode before downloading model files — GPU requires the FP32 model.
        // DAKERA_USE_GPU=1 overrides config.use_gpu (env var is the production knob).
        let use_gpu = std::env::var("DAKERA_USE_GPU")
            .map(|v| v == "1")
            .unwrap_or(config.use_gpu);
        if use_gpu {
            info!("CUDA execution provider enabled — using FP32 model (DAKERA_USE_GPU=1)");
        }

        info!(
            "Initializing ONNX embedding engine with model: {}",
            config.model
        );

        // Download tokenizer and ONNX model files (FP32 for GPU, INT8 for CPU)
        let (tokenizer_path, onnx_path) = Self::download_model_files(&config, use_gpu).await?;

        // Load tokenizer
        info!("Loading tokenizer from {:?}", tokenizer_path);
        let tokenizer = Tokenizer::from_file(&tokenizer_path)
            .map_err(|e| InferenceError::TokenizationError(e.to_string()))?;

        // Build ONNX session pool — N independent sessions to serve concurrent callers.
        // Each session has its own ORT context so pool members never block each other.
        info!("Loading ONNX model from {:?}", onnx_path);
        let num_threads = config.num_threads.unwrap_or(4);
        let pool_size = config.session_pool_size.max(1);
        let onnx_path_clone = onnx_path.clone();

        let sessions: Vec<Arc<Mutex<Session>>> =
            tokio::task::spawn_blocking(move || -> Result<Vec<Arc<Mutex<Session>>>> {
                (0..pool_size)
                    .map(|_| {
                        let builder = Session::builder()
                            .map_err(|e| InferenceError::ModelLoadError(e.to_string()))?
                            .with_optimization_level(GraphOptimizationLevel::Level3)
                            .map_err(|e| InferenceError::ModelLoadError(e.to_string()))?
                            .with_intra_threads(num_threads)
                            .map_err(|e| InferenceError::ModelLoadError(e.to_string()))?;

                        let mut builder = if use_gpu {
                            builder
                                .with_execution_providers(
                                    [CUDAExecutionProvider::default().build()],
                                )
                                .map_err(|e| InferenceError::ModelLoadError(e.to_string()))?
                        } else {
                            builder
                        };

                        let s = builder
                            .commit_from_file(&onnx_path_clone)
                            .map_err(|e| InferenceError::ModelLoadError(e.to_string()))?;
                        Ok(Arc::new(Mutex::new(s)))
                    })
                    .collect()
            })
            .await
            .map_err(|e| {
                InferenceError::ModelLoadError(format!("Session pool init panicked: {}", e))
            })??;

        let dimension = config.model.dimension();
        let processor = Arc::new(BatchProcessor::new(
            tokenizer,
            config.model,
            config.max_batch_size,
        ));

        info!(
            "ONNX embedding engine ready: model={}, dimension={}, threads={}, pool={}",
            config.model, dimension, num_threads, pool_size
        );

        Ok(Self {
            sessions,
            next_session: AtomicUsize::new(0),
            processor,
            config,
            dimension,
        })
    }

    /// Resolve tokenizer and ONNX model files, downloading from HuggingFace if needed.
    ///
    /// - `tokenizer.json` — from the original model repo (sentence-transformers, BAAI, intfloat)
    /// - GPU: `onnx/model.onnx` (FP32) — no Memcpy round-trips, CUDA EP runs entirely on-device
    /// - CPU: `onnx/model_quantized.onnx` (INT8) — fastest for CPU-only deployments
    #[instrument(skip_all, fields(model = %config.model))]
    async fn download_model_files(
        config: &ModelConfig,
        use_gpu: bool,
    ) -> Result<(PathBuf, PathBuf)> {
        let model_id = config.model.model_id();
        let onnx_repo_id = config.model.onnx_repo_id();
        let onnx_filename = if use_gpu {
            config.model.onnx_filename_gpu()
        } else {
            config.model.onnx_filename()
        };

        info!(
            "Resolving model files: tokenizer={}, onnx={}@{}",
            model_id, onnx_filename, onnx_repo_id
        );

        let tokenizer_cache_dir = Self::model_cache_dir(model_id)?;
        let onnx_cache_dir = Self::model_cache_dir(onnx_repo_id)?;

        // ONNX sub-directory mirrors the path within the repo (e.g. "onnx/")
        let onnx_subdir = onnx_cache_dir.join("onnx");
        std::fs::create_dir_all(&onnx_subdir)?;

        let local_tokenizer = tokenizer_cache_dir.join("tokenizer.json");
        // onnx_filename is e.g. "onnx/model_quantized.onnx" or "onnx/model.onnx" — use the basename
        let onnx_basename = Path::new(onnx_filename)
            .file_name()
            .and_then(|s| s.to_str())
            .unwrap_or("model_quantized.onnx");
        let local_onnx = onnx_subdir.join(onnx_basename);

        // Download missing files in a blocking thread
        let tokenizer_needs_download = !local_tokenizer.exists();

        // GPU FP32 model (BGE-large) is 1.27 GB. The old 500 MB download limit silently
        // truncated it, leaving a corrupt ONNX file that ORT rejects with "Protobuf parsing
        // failed" (DAK-5976). If the cached GPU model is ≤500 MB, it's the truncated artifact
        // — delete it so download_hf_file fetches the complete file.
        if use_gpu && local_onnx.exists() {
            let cached_size = local_onnx.metadata().map(|m| m.len()).unwrap_or(0);
            if cached_size <= 500_000_000 {
                warn!(
                    "Cached GPU ONNX at {:?} is {} bytes (≤500 MB) — likely truncated by old \
                     download limit. Deleting for complete re-download.",
                    local_onnx, cached_size
                );
                let _ = std::fs::remove_file(&local_onnx);
            }
        }
        let onnx_needs_download = !local_onnx.exists();

        if tokenizer_needs_download || onnx_needs_download {
            let model_id_owned = model_id.to_string();
            let onnx_repo_id_owned = onnx_repo_id.to_string();
            let onnx_filename_owned = onnx_filename.to_string();
            let tokenizer_cache = tokenizer_cache_dir.clone();
            let onnx_cache = onnx_cache_dir.clone();

            tokio::task::spawn_blocking(move || {
                if !tokenizer_cache.join("tokenizer.json").exists() {
                    Self::download_hf_file(&model_id_owned, "tokenizer.json", &tokenizer_cache)
                        .map_err(|e| {
                            InferenceError::HubError(format!("Failed to download tokenizer: {}", e))
                        })?;
                }
                if !onnx_cache.join(&onnx_filename_owned).exists() {
                    Self::download_hf_file(&onnx_repo_id_owned, &onnx_filename_owned, &onnx_cache)
                        .map_err(|e| {
                            InferenceError::HubError(format!(
                                "Failed to download ONNX model: {}",
                                e
                            ))
                        })?;
                }
                Ok::<_, InferenceError>(())
            })
            .await
            .map_err(|e| InferenceError::HubError(format!("Download task panicked: {}", e)))??;
        } else {
            info!("All model files found in local cache");
        }

        // Re-derive paths (cache dir / onnx / basename)
        let final_onnx = onnx_cache_dir.join(onnx_filename);

        info!(
            "Model files ready: tokenizer={:?}, onnx={:?}",
            local_tokenizer, final_onnx
        );
        Ok((local_tokenizer, final_onnx))
    }

    /// Get or create the local model cache directory.
    fn model_cache_dir(model_id: &str) -> Result<PathBuf> {
        let base = std::env::var("HF_HOME")
            .map(PathBuf::from)
            .unwrap_or_else(|_| {
                let home = std::env::var("HOME").unwrap_or_else(|_| {
                    warn!("HOME environment variable not set, using /tmp for model cache");
                    "/tmp".to_string()
                });
                PathBuf::from(home).join(".cache").join("huggingface")
            });
        let dir = base.join("dakera").join(model_id.replace('/', "--"));
        std::fs::create_dir_all(&dir)?;
        Ok(dir)
    }

    /// Download a single file from HuggingFace using ureq (sync, for spawn_blocking).
    ///
    /// Handles relative Location headers that ureq 2.x cannot resolve automatically.
    ///
    /// Public alias for use by other inference modules (e.g. GLiNER NER engine).
    pub fn download_hf_file_pub(
        model_id: &str,
        filename: &str,
        cache_dir: &Path,
    ) -> std::result::Result<PathBuf, String> {
        Self::download_hf_file(model_id, filename, cache_dir)
    }

    fn download_hf_file(
        model_id: &str,
        filename: &str,
        cache_dir: &Path,
    ) -> std::result::Result<PathBuf, String> {
        // The file may be nested (e.g. "onnx/model_quantized.onnx")
        let file_path = cache_dir.join(filename);
        if file_path.exists() {
            info!("Cached: {}/{}", model_id, filename);
            return Ok(file_path);
        }

        // Ensure parent directory exists (for "onnx/model_quantized.onnx")
        if let Some(parent) = file_path.parent() {
            std::fs::create_dir_all(parent)
                .map_err(|e| format!("Failed to create directory {:?}: {}", parent, e))?;
        }

        let url = format!(
            "https://huggingface.co/{}/resolve/main/{}",
            model_id, filename
        );
        info!("Downloading: {}", url);

        // Read HuggingFace token from env (required for gated models).
        let hf_token = std::env::var("HF_TOKEN")
            .or_else(|_| std::env::var("HUGGING_FACE_HUB_TOKEN"))
            .ok();
        if hf_token.is_some() {
            info!("Using HuggingFace auth token for download");
        }

        // Disable automatic redirects so we can resolve relative Location headers ourselves.
        let agent = ureq::AgentBuilder::new()
            .redirects(0)
            .timeout(std::time::Duration::from_secs(300))
            .build();

        let mut current_url = url.clone();
        let mut redirects = 0;
        let max_redirects = 10;

        let response = loop {
            let mut req = agent.get(&current_url);
            if let Some(ref token) = hf_token {
                req = req.set("Authorization", &format!("Bearer {}", token));
            }
            let resp = req.call();

            let r = match resp {
                Ok(r) => r,
                Err(ureq::Error::Status(_status, r)) => r,
                Err(e) => return Err(format!("{}: {}", filename, e)),
            };

            let status = r.status();
            if (200..300).contains(&status) {
                break r;
            } else if (300..400).contains(&status) {
                redirects += 1;
                if redirects > max_redirects {
                    return Err(format!("{}: too many redirects", filename));
                }
                let location = r
                    .header("location")
                    .ok_or_else(|| format!("{}: redirect without Location header", filename))?
                    .to_string();

                // Resolve relative redirects against the current URL's origin
                current_url = if location.starts_with('/') {
                    let parsed = url::Url::parse(&current_url)
                        .map_err(|e| format!("{}: bad URL {}: {}", filename, current_url, e))?;
                    let host = parsed.host_str().ok_or_else(|| {
                        format!("{}: redirect URL missing host: {}", filename, current_url)
                    })?;
                    format!("{}://{}{}", parsed.scheme(), host, location)
                } else {
                    location
                };
                info!("Redirect {} → {}", redirects, current_url);
            } else {
                return Err(format!("{}: HTTP {}", filename, status));
            }
        };

        // Capture expected file size before consuming the response body.
        // HuggingFace CDN reports LFS blob sizes via x-linked-size; standard
        // HTTP servers use content-length. Either may be absent (chunked transfer).
        let expected_bytes: Option<u64> = response
            .header("x-linked-size")
            .or_else(|| response.header("content-length"))
            .and_then(|v| v.parse::<u64>().ok());

        // 2 GiB limit — BGE-large FP32 is 1.27 GB; the old 500 MB limit silently
        // truncated it, producing a corrupt ONNX that ORT rejects with
        // "Protobuf parsing failed" (DAK-5976).
        let mut bytes = Vec::new();
        response
            .into_reader()
            .take(2_147_483_648)
            .read_to_end(&mut bytes)
            .map_err(|e| format!("Failed to read {}: {}", filename, e))?;

        // Fail fast if the received payload is shorter than the declared size.
        // This prevents writing a corrupt file that would pass the cache-hit check
        // on the next boot but fail when ORT tries to parse it.
        if let Some(expected) = expected_bytes {
            let actual = bytes.len() as u64;
            if actual < expected {
                return Err(format!(
                    "{}: download incomplete — received {} of {} bytes. \
                     File may exceed 2 GiB or the connection was interrupted.",
                    filename, actual, expected
                ));
            }
        }

        std::fs::write(&file_path, &bytes)
            .map_err(|e| format!("Failed to write {}: {}", filename, e))?;

        info!("Downloaded {} ({} bytes)", filename, bytes.len());
        Ok(file_path)
    }

    /// Get the embedding dimension for the loaded model.
    pub fn dimension(&self) -> usize {
        self.dimension
    }

    /// Get the model being used.
    pub fn model(&self) -> EmbeddingModel {
        self.config.model
    }

    /// Get the number of parallel ONNX sessions in the pool.
    pub fn pool_size(&self) -> usize {
        self.sessions.len()
    }

    /// Embed a single query text.
    ///
    /// For models like E5, this automatically applies the query prefix.
    #[instrument(skip(self, text), fields(text_len = text.len()))]
    pub async fn embed_query(&self, text: &str) -> Result<Vec<f32>> {
        let texts = vec![text.to_string()];
        let prepared = self.processor.prepare_texts(&texts, true);
        let embeddings = self.embed_batch_internal(&prepared).await?;
        embeddings.into_iter().next().ok_or_else(|| {
            InferenceError::InferenceError("No embedding returned for query".to_string())
        })
    }

    /// Embed multiple query texts.
    ///
    /// For models like E5, this automatically applies the query prefix.
    #[instrument(skip(self, texts), fields(count = texts.len()))]
    pub async fn embed_queries(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
        let prepared = self.processor.prepare_texts(texts, true);
        self.embed_batch_internal(&prepared).await
    }

    /// Embed a single document/passage.
    ///
    /// For models like E5, this automatically applies the document prefix.
    #[instrument(skip(self, text), fields(text_len = text.len()))]
    pub async fn embed_document(&self, text: &str) -> Result<Vec<f32>> {
        let texts = vec![text.to_string()];
        let prepared = self.processor.prepare_texts(&texts, false);
        let embeddings = self.embed_batch_internal(&prepared).await?;
        embeddings.into_iter().next().ok_or_else(|| {
            InferenceError::InferenceError("No embedding returned for document".to_string())
        })
    }

    /// Embed multiple documents/passages.
    ///
    /// For models like E5, this automatically applies the document prefix.
    #[instrument(skip(self, texts), fields(count = texts.len()))]
    pub async fn embed_documents(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
        let prepared = self.processor.prepare_texts(texts, false);
        self.embed_batch_internal(&prepared).await
    }

    /// Embed texts without any prefix (raw embedding).
    #[instrument(skip(self, texts), fields(count = texts.len()))]
    pub async fn embed_raw(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
        self.embed_batch_internal(texts).await
    }

    /// Internal batch embedding implementation.
    ///
    /// Splits `texts` into sub-batches (≤ `batch_size` each) and distributes them
    /// across the session pool via round-robin.  On GPU BFCArena OOM the batch size
    /// is halved and the full input is retried (up to 3 halvings; DAK-6002).  This
    /// handles LME-style long documents whose attention tensor (batch × heads × seq²)
    /// exceeds the VRAM headroom at the configured batch size but fits at batch/2.
    async fn embed_batch_internal(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
        if texts.is_empty() {
            return Ok(vec![]);
        }

        let pool_len = self.sessions.len();
        let normalize = self.config.model.normalize_embeddings();
        // Round-robin starting index: each concurrent caller gets a different slot so
        // concurrent requests don't all contend on sessions[0].
        let start_idx = self.next_session.fetch_add(1, Ordering::Relaxed);

        let mut batch_size = self.config.max_batch_size.max(1);

        // Up to 3 halvings: 32 → 16 → 8 → 4 → hard fail.
        for attempt in 0_u32..=3 {
            let batches: Vec<Vec<String>> = texts.chunks(batch_size).map(|b| b.to_vec()).collect();

            // Spawn all sub-batches concurrently; preserve insertion order for reassembly.
            let mut handles = Vec::with_capacity(batches.len());
            for (i, batch_owned) in batches.into_iter().enumerate() {
                let session = Arc::clone(&self.sessions[(start_idx + i) % pool_len]);
                let processor = Arc::clone(&self.processor);
                handles.push(tokio::task::spawn_blocking(move || {
                    let mut session_guard = session.lock();
                    Self::process_batch_blocking(
                        &batch_owned,
                        &mut session_guard,
                        &processor,
                        normalize,
                    )
                }));
            }

            let mut all_embeddings: Vec<Vec<f32>> = Vec::with_capacity(texts.len());
            let mut oom: Option<InferenceError> = None;

            for handle in handles {
                match handle.await {
                    Err(panic_err) => {
                        return Err(InferenceError::InferenceError(format!(
                            "Inference task panicked: {panic_err}"
                        )));
                    }
                    Ok(Err(e)) => {
                        if attempt < 3 && Self::is_gpu_oom(&e) {
                            // First OOM wins; remaining in-flight tasks detach harmlessly.
                            oom = Some(e);
                            break;
                        }
                        return Err(e);
                    }
                    Ok(Ok(batch_embs)) => {
                        all_embeddings.extend(batch_embs);
                    }
                }
            }

            if oom.is_some() {
                let next_batch = (batch_size / 2).max(1);
                warn!(
                    "ONNX allocator OOM (attempt {}/3) — retrying with batch_size {} → {}",
                    attempt + 1,
                    batch_size,
                    next_batch,
                );
                batch_size = next_batch;
                continue;
            }

            return Ok(all_embeddings);
        }

        Err(InferenceError::InferenceError(format!(
            "ONNX inference failed: GPU/CPU allocator OOM persists after 3 \
             batch-halving attempts (final batch_size={batch_size})"
        )))
    }

    /// Returns `true` when the error message indicates an allocator OOM that can be
    /// recovered by reducing the inference batch size (DAK-6002).
    ///
    /// Covers CUDA BFCArena (`BFCArena`, `Failed to allocate memory … buffer of size`)
    /// and CPU arena OOM patterns seen in ONNX Runtime 2.x.
    fn is_gpu_oom(err: &InferenceError) -> bool {
        let msg = err.to_string();
        msg.contains("BFCArena")
            || msg.contains("Failed to allocate memory")
            || msg.contains("CUDA_OUT_OF_MEMORY")
            || msg.contains("CUDA out of memory")
            || (msg.contains("allocate") && msg.contains("buffer of size"))
    }

    /// Process a single batch: tokenize → ORT session → mean pool → normalize.
    ///
    /// Designed to run inside `spawn_blocking` (takes no `&self`).
    fn process_batch_blocking(
        texts: &[String],
        session: &mut Session,
        processor: &BatchProcessor,
        normalize: bool,
    ) -> Result<Vec<Vec<f32>>> {
        // Tokenize
        let prepared = processor.tokenize_batch(texts)?;
        let batch_size = prepared.batch_size;
        let seq_len = prepared.seq_len;

        // Keep a copy of attention_mask for mean_pooling (consumed by Tensor below)
        let attention_mask_flat = prepared.attention_mask.clone();

        // Build ORT tensors — from_array requires (shape, Vec<T>) in ort rc.12
        let input_ids_tensor =
            Tensor::<i64>::from_array(([batch_size, seq_len], prepared.input_ids))
                .map_err(|e| InferenceError::InferenceError(e.to_string()))?;
        let attention_mask_tensor =
            Tensor::<i64>::from_array(([batch_size, seq_len], prepared.attention_mask))
                .map_err(|e| InferenceError::InferenceError(e.to_string()))?;
        let token_type_ids_tensor =
            Tensor::<i64>::from_array(([batch_size, seq_len], prepared.token_type_ids))
                .map_err(|e| InferenceError::InferenceError(e.to_string()))?;

        // Run ONNX session
        let outputs = session
            .run(inputs![
                "input_ids" => input_ids_tensor,
                "attention_mask" => attention_mask_tensor,
                "token_type_ids" => token_type_ids_tensor
            ])
            .map_err(|e: ort::Error| InferenceError::InferenceError(e.to_string()))?;

        // Extract last_hidden_state: shape [batch, seq_len, hidden_size]
        // ort rc.12: try_extract_tensor returns (&Shape, &[T])
        // Shape derefs to [i64], so index directly.
        let (ort_shape, lhs_slice) = outputs[0]
            .try_extract_tensor::<f32>()
            .map_err(|e| InferenceError::InferenceError(e.to_string()))?;

        if ort_shape.len() != 3 {
            return Err(InferenceError::InferenceError(format!(
                "Expected 3D last_hidden_state, got {} dims",
                ort_shape.len()
            )));
        }
        let hidden_size = ort_shape[2] as usize;

        // Apply mean pooling using the saved attention mask copy
        let mut embeddings = mean_pooling(
            lhs_slice,
            batch_size,
            seq_len,
            hidden_size,
            &attention_mask_flat,
        );

        // L2 normalize if configured
        if normalize {
            normalize_embeddings(&mut embeddings);
        }

        debug!(
            "Generated {} embeddings of dimension {}",
            embeddings.len(),
            embeddings.first().map(|e| e.len()).unwrap_or(0)
        );

        Ok(embeddings)
    }

    /// Estimate the time to embed a batch of texts (in milliseconds).
    pub fn estimate_time_ms(&self, text_count: usize, avg_text_len: usize) -> f64 {
        // Rough estimation based on model speed and text length (CPU path)
        let tokens_per_text =
            (avg_text_len as f64 / 4.0).min(self.config.model.max_seq_length() as f64);
        let total_tokens = tokens_per_text * text_count as f64;
        let tokens_per_second = self.config.model.tokens_per_second_cpu() as f64;
        (total_tokens / tokens_per_second) * 1000.0
    }
}

impl std::fmt::Debug for EmbeddingEngine {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("EmbeddingEngine")
            .field("model", &self.config.model)
            .field("dimension", &self.dimension)
            .field("max_batch_size", &self.config.max_batch_size)
            .field("session_pool_size", &self.sessions.len())
            .finish()
    }
}

/// Builder for creating an EmbeddingEngine with fluent API.
pub struct EmbeddingEngineBuilder {
    config: ModelConfig,
}

impl EmbeddingEngineBuilder {
    /// Create a new builder with default configuration.
    pub fn new() -> Self {
        Self {
            config: ModelConfig::default(),
        }
    }

    /// Set the embedding model to use.
    pub fn model(mut self, model: EmbeddingModel) -> Self {
        self.config.model = model;
        self
    }

    /// Set the cache directory for model files.
    pub fn cache_dir(mut self, dir: impl Into<String>) -> Self {
        self.config.cache_dir = Some(dir.into());
        self
    }

    /// Set the maximum batch size.
    pub fn max_batch_size(mut self, size: usize) -> Self {
        self.config.max_batch_size = size;
        self
    }

    /// Enable GPU acceleration (reserved for future use; ORT selects the execution provider).
    pub fn use_gpu(mut self, enable: bool) -> Self {
        self.config.use_gpu = enable;
        self
    }

    /// Set the number of intra-op CPU threads for ORT inference.
    pub fn num_threads(mut self, threads: usize) -> Self {
        self.config.num_threads = Some(threads);
        self
    }

    /// Set the number of parallel ONNX sessions in the pool.
    pub fn session_pool_size(mut self, size: usize) -> Self {
        self.config.session_pool_size = size.max(1);
        self
    }

    /// Build the embedding engine.
    pub async fn build(self) -> Result<EmbeddingEngine> {
        EmbeddingEngine::new(self.config).await
    }
}

impl Default for EmbeddingEngineBuilder {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_estimate_time() {
        let config = ModelConfig::new(EmbeddingModel::MiniLM);
        let tokens_per_second = config.model.tokens_per_second_cpu() as f64;
        assert!(tokens_per_second > 0.0);
    }

    #[test]
    fn test_builder() {
        let builder = EmbeddingEngineBuilder::new()
            .model(EmbeddingModel::BgeSmall)
            .max_batch_size(64)
            .use_gpu(false);

        assert_eq!(builder.config.model, EmbeddingModel::BgeSmall);
        assert_eq!(builder.config.max_batch_size, 64);
        assert!(!builder.config.use_gpu);
    }

    // ── model_cache_dir ──────────────────────────────────────────────────────

    /// Ensure `model_cache_dir` respects `HF_HOME` when set.
    ///
    /// Uses a process-level mutex because `std::env::set_var` is not thread-safe.
    #[test]
    fn test_model_cache_dir_with_hf_home() {
        use std::sync::Mutex;
        static ENV_LOCK: Mutex<()> = Mutex::new(());
        let _guard = ENV_LOCK.lock().unwrap();

        let tmp = std::env::temp_dir().join("dakera_test_hf_home");
        std::env::set_var("HF_HOME", &tmp);
        let result = EmbeddingEngine::model_cache_dir("org/my-model");
        std::env::remove_var("HF_HOME");

        let path = result.unwrap();
        assert!(
            path.starts_with(&tmp),
            "expected path under {tmp:?}, got {path:?}"
        );
        assert!(
            path.to_str().unwrap().contains("org--my-model"),
            "model_id separator not applied: {path:?}"
        );
    }

    #[test]
    fn test_model_cache_dir_contains_dakera_subdir() {
        let path =
            EmbeddingEngine::model_cache_dir("sentence-transformers/all-MiniLM-L6-v2").unwrap();
        let s = path.to_str().unwrap();
        assert!(s.contains("dakera"), "expected 'dakera' in path: {s}");
        assert!(
            s.contains("sentence-transformers--all-MiniLM-L6-v2"),
            "expected transformed model id in path: {s}"
        );
    }

    #[test]
    fn test_model_cache_dir_creates_directory() {
        let dir = EmbeddingEngine::model_cache_dir("test/cache-dir-creation-probe").unwrap();
        assert!(dir.exists(), "model_cache_dir should create the directory");
    }

    // ── download_hf_file (cached early-return path) ──────────────────────────

    #[test]
    fn test_download_hf_file_returns_path_when_already_cached() {
        let tmp = std::env::temp_dir().join("dakera_test_cached_file");
        std::fs::create_dir_all(&tmp).unwrap();
        let file_path = tmp.join("config.json");
        std::fs::write(&file_path, b"{}").unwrap();

        let result = EmbeddingEngine::download_hf_file("test/model", "config.json", &tmp);
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), file_path);
    }

    #[test]
    fn test_download_hf_file_returns_correct_path_for_cached_onnx() {
        let tmp = std::env::temp_dir().join("dakera_test_cached_onnx");
        let onnx_dir = tmp.join("onnx");
        std::fs::create_dir_all(&onnx_dir).unwrap();
        let onnx_path = onnx_dir.join("model_quantized.onnx");
        std::fs::write(&onnx_path, b"fake_onnx_model").unwrap();

        // Filename includes the subdirectory path
        let result = EmbeddingEngine::download_hf_file(
            "Xenova/all-MiniLM-L6-v2",
            "onnx/model_quantized.onnx",
            &tmp,
        );
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), onnx_path);
    }

    // ── EmbeddingEngineBuilder ───────────────────────────────────────────────

    #[test]
    fn test_builder_default_impl() {
        let b1 = EmbeddingEngineBuilder::new();
        let b2 = EmbeddingEngineBuilder::default();
        assert_eq!(b1.config.model, b2.config.model);
        assert_eq!(b1.config.max_batch_size, b2.config.max_batch_size);
    }

    #[test]
    fn test_builder_model_field() {
        let builder = EmbeddingEngineBuilder::new().model(EmbeddingModel::E5Small);
        assert_eq!(builder.config.model, EmbeddingModel::E5Small);
    }

    #[test]
    fn test_builder_cache_dir() {
        let builder = EmbeddingEngineBuilder::new().cache_dir("/tmp/my-models");
        assert_eq!(builder.config.cache_dir, Some("/tmp/my-models".to_string()));
    }

    #[test]
    fn test_builder_max_batch_size() {
        let builder = EmbeddingEngineBuilder::new().max_batch_size(128);
        assert_eq!(builder.config.max_batch_size, 128);
    }

    #[test]
    fn test_builder_use_gpu_true() {
        let builder = EmbeddingEngineBuilder::new().use_gpu(true);
        assert!(builder.config.use_gpu);
    }

    #[test]
    fn test_builder_use_gpu_false() {
        let builder = EmbeddingEngineBuilder::new().use_gpu(false);
        assert!(!builder.config.use_gpu);
    }

    #[test]
    fn test_builder_num_threads() {
        let builder = EmbeddingEngineBuilder::new().num_threads(4);
        assert_eq!(builder.config.num_threads, Some(4));
    }

    #[test]
    fn test_builder_chain_all_fields() {
        let builder = EmbeddingEngineBuilder::new()
            .model(EmbeddingModel::BgeSmall)
            .cache_dir("/cache")
            .max_batch_size(16)
            .use_gpu(false)
            .num_threads(2);

        assert_eq!(builder.config.model, EmbeddingModel::BgeSmall);
        assert_eq!(builder.config.cache_dir, Some("/cache".to_string()));
        assert_eq!(builder.config.max_batch_size, 16);
        assert!(!builder.config.use_gpu);
        assert_eq!(builder.config.num_threads, Some(2));
    }

    // ── estimate_time_ms ─────────────────────────────────────────────────────

    #[test]
    fn test_estimate_time_zero_count() {
        let tps = EmbeddingModel::MiniLM.tokens_per_second_cpu() as f64;
        let estimate = (0.0 / tps) * 1000.0;
        assert_eq!(estimate, 0.0);
    }

    #[test]
    fn test_estimate_time_formula_cpu() {
        // texts=10, avg_len=100 → tokens_per_text = min(25, 256) = 25
        // total_tokens = 250; tps = 5000; time = (250/5000)*1000 = 50ms
        let model = EmbeddingModel::MiniLM;
        let tokens_per_text = (100.0f64 / 4.0).min(model.max_seq_length() as f64);
        let total_tokens = tokens_per_text * 10.0;
        let estimate = (total_tokens / model.tokens_per_second_cpu() as f64) * 1000.0;
        assert!(
            (estimate - 50.0).abs() < 1e-6,
            "expected 50.0ms, got {estimate}"
        );
    }

    #[test]
    fn test_estimate_time_capped_at_max_seq_length() {
        let model = EmbeddingModel::MiniLM;
        let avg_len = 100_000;
        let tokens_per_text = (avg_len as f64 / 4.0).min(model.max_seq_length() as f64);
        assert_eq!(tokens_per_text, 256.0);
    }

    // ── ModelConfig API ───────────────────────────────────────────────────────

    #[test]
    fn test_model_config_new() {
        let cfg = ModelConfig::new(EmbeddingModel::BgeSmall);
        assert_eq!(cfg.model, EmbeddingModel::BgeSmall);
        assert_eq!(cfg.max_batch_size, 32);
        assert!(!cfg.use_gpu);
        assert!(cfg.cache_dir.is_none());
        assert!(cfg.num_threads.is_none());
    }

    #[test]
    fn test_model_config_default() {
        let cfg = ModelConfig::default();
        assert_eq!(cfg.model, EmbeddingModel::BgeLarge);
        assert_eq!(cfg.max_batch_size, 32);
        assert!(!cfg.use_gpu);
    }

    #[test]
    fn test_model_config_with_cache_dir() {
        let cfg = ModelConfig::new(EmbeddingModel::MiniLM).with_cache_dir("/tmp/models");
        assert_eq!(cfg.cache_dir, Some("/tmp/models".to_string()));
    }

    #[test]
    fn test_model_config_with_max_batch_size() {
        let cfg = ModelConfig::new(EmbeddingModel::MiniLM).with_max_batch_size(64);
        assert_eq!(cfg.max_batch_size, 64);
    }

    #[test]
    fn test_model_config_with_gpu() {
        let cfg = ModelConfig::new(EmbeddingModel::MiniLM).with_gpu(true);
        assert!(cfg.use_gpu);
    }

    #[test]
    fn test_model_config_with_num_threads() {
        let cfg = ModelConfig::new(EmbeddingModel::MiniLM).with_num_threads(8);
        assert_eq!(cfg.num_threads, Some(8));
    }

    #[test]
    fn test_model_config_chained_builder() {
        let cfg = ModelConfig::new(EmbeddingModel::E5Small)
            .with_cache_dir("/cache")
            .with_max_batch_size(16)
            .with_gpu(false)
            .with_num_threads(4);
        assert_eq!(cfg.model, EmbeddingModel::E5Small);
        assert_eq!(cfg.cache_dir, Some("/cache".to_string()));
        assert_eq!(cfg.max_batch_size, 16);
        assert!(!cfg.use_gpu);
        assert_eq!(cfg.num_threads, Some(4));
    }

    // ── model_cache_dir edge cases ────────────────────────────────────────────

    /// Test `model_cache_dir` when HOME is not set — should fall back to /tmp.
    #[test]
    fn test_model_cache_dir_no_home_fallback() {
        use std::sync::Mutex;
        static ENV_LOCK: Mutex<()> = Mutex::new(());
        let _guard = ENV_LOCK.lock().unwrap();

        // Remove HOME and HF_HOME so we hit the /tmp fallback
        let saved_home = std::env::var("HOME").ok();
        let saved_hf = std::env::var("HF_HOME").ok();
        unsafe {
            std::env::remove_var("HOME");
            std::env::remove_var("HF_HOME");
        }

        let result = EmbeddingEngine::model_cache_dir("test/fallback-model");

        // Restore env
        if let Some(h) = saved_home {
            unsafe { std::env::set_var("HOME", h) };
        }
        if let Some(h) = saved_hf {
            unsafe { std::env::set_var("HF_HOME", h) };
        }

        let path = result.unwrap();
        // Should be under /tmp since HOME was unset
        assert!(
            path.starts_with("/tmp"),
            "expected path under /tmp, got {path:?}"
        );
    }

    #[test]
    fn test_model_cache_dir_deep_model_id() {
        let path = EmbeddingEngine::model_cache_dir("org/sub/model-name-with-dashes").unwrap();
        let s = path.to_str().unwrap();
        // All slashes replaced with double-dash
        assert!(
            s.contains("org--sub--model-name-with-dashes"),
            "expected transformed path, got: {s}"
        );
    }

    #[test]
    fn test_model_cache_dir_minilm_model_id() {
        let path = EmbeddingEngine::model_cache_dir(EmbeddingModel::MiniLM.model_id()).unwrap();
        let s = path.to_str().unwrap();
        assert!(s.contains("sentence-transformers--all-MiniLM-L6-v2"));
    }

    #[test]
    fn test_model_cache_dir_bge_model_id() {
        let path = EmbeddingEngine::model_cache_dir(EmbeddingModel::BgeSmall.model_id()).unwrap();
        let s = path.to_str().unwrap();
        assert!(s.contains("BAAI--bge-small-en-v1.5"));
    }

    #[test]
    fn test_model_cache_dir_e5_model_id() {
        let path = EmbeddingEngine::model_cache_dir(EmbeddingModel::E5Small.model_id()).unwrap();
        let s = path.to_str().unwrap();
        assert!(s.contains("intfloat--e5-small-v2"));
    }

    // ── download_hf_file additional cache-hit variations ─────────────────────

    #[test]
    fn test_download_hf_file_pytorch_bin_cached() {
        let tmp = std::env::temp_dir().join("dakera_test_pytorch_bin");
        std::fs::create_dir_all(&tmp).unwrap();
        let model_path = tmp.join("pytorch_model.bin");
        std::fs::write(&model_path, b"fake_pytorch_weights").unwrap();

        let result = EmbeddingEngine::download_hf_file("test/model", "pytorch_model.bin", &tmp);
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), model_path);
    }

    #[test]
    fn test_download_hf_file_tokenizer_cached() {
        let tmp = std::env::temp_dir().join("dakera_test_tokenizer_cached");
        std::fs::create_dir_all(&tmp).unwrap();
        let tok_path = tmp.join("tokenizer.json");
        std::fs::write(&tok_path, br#"{"version":"1.0"}"#).unwrap();

        let result = EmbeddingEngine::download_hf_file("test/model", "tokenizer.json", &tmp);
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), tok_path);
    }

    #[test]
    fn test_download_hf_file_config_json_cached() {
        let tmp = std::env::temp_dir().join("dakera_test_config_cached");
        std::fs::create_dir_all(&tmp).unwrap();
        let cfg_path = tmp.join("config.json");
        std::fs::write(&cfg_path, b"{}").unwrap();

        let result = EmbeddingEngine::download_hf_file("test/model", "config.json", &tmp);
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), cfg_path);
    }

    // ── EmbeddingEngine::new() failure path via fake local cache ─────────────

    /// Tests the code path through `download_model_files` (local Dakera cache hit)
    /// and into `new()` — which then fails trying to load the tokenizer from a
    /// fake file. No network access required; fake files are pre-seeded.
    #[tokio::test]
    #[allow(clippy::await_holding_lock)]
    async fn test_new_fails_with_invalid_tokenizer_json() {
        use std::sync::Mutex;
        static ENV_LOCK: Mutex<()> = Mutex::new(());
        let _guard = ENV_LOCK.lock().unwrap();

        // Set up a fake Dakera model cache so download_model_files finds our files
        let tmp = std::env::temp_dir().join("dakera_test_engine_new_fail_tok");
        let model_dir = tmp
            .join("dakera")
            .join("sentence-transformers--all-MiniLM-L6-v2");
        std::fs::create_dir_all(&model_dir).unwrap();
        // Valid-looking model weights placeholder (candle will fail on this, which is fine)
        std::fs::write(model_dir.join("model.safetensors"), b"not_real_weights").unwrap();
        // Invalid tokenizer.json — will cause TokenizationError in new()
        std::fs::write(model_dir.join("tokenizer.json"), b"NOT_VALID_JSON").unwrap();
        std::fs::write(model_dir.join("config.json"), b"{}").unwrap();

        unsafe { std::env::set_var("HF_HOME", &tmp) };

        let config = ModelConfig::new(EmbeddingModel::MiniLM);
        let result = EmbeddingEngine::new(config).await;

        unsafe { std::env::remove_var("HF_HOME") };

        // Must fail — tokenizer.json is invalid JSON
        assert!(
            result.is_err(),
            "expected Err from new() with invalid tokenizer, got Ok"
        );
    }

    // ── EmbeddingEngineBuilder additional coverage ────────────────────────────

    #[test]
    fn test_builder_with_all_models() {
        for model in [
            EmbeddingModel::MiniLM,
            EmbeddingModel::BgeSmall,
            EmbeddingModel::E5Small,
        ] {
            let builder = EmbeddingEngineBuilder::new().model(model);
            assert_eq!(builder.config.model, model);
        }
    }

    #[test]
    fn test_builder_max_batch_size_one() {
        let builder = EmbeddingEngineBuilder::new().max_batch_size(1);
        assert_eq!(builder.config.max_batch_size, 1);
    }

    #[test]
    fn test_builder_num_threads_zero() {
        let builder = EmbeddingEngineBuilder::new().num_threads(0);
        assert_eq!(builder.config.num_threads, Some(0));
    }

    // ── EmbeddingEngine::new() / getters when model is cached (best-effort) ──

    /// If the embedding model is already cached on this machine, exercise the
    /// full `new()` path and test getters. On machines without a cached model
    /// the test passes silently — it is intentionally non-gating.
    #[tokio::test]
    async fn test_engine_getters_when_model_cached() {
        let config = ModelConfig::new(EmbeddingModel::MiniLM);
        match EmbeddingEngine::new(config).await {
            Ok(engine) => {
                assert_eq!(engine.dimension(), EmbeddingModel::MiniLM.dimension());
                assert_eq!(engine.model(), EmbeddingModel::MiniLM);
                // Device should be CPU in test environments (device() removed in CE-3 ONNX migration)
                // Debug impl should not panic
                let _ = format!("{:?}", engine);
                // estimate_time_ms should return a non-negative value
                let ms = engine.estimate_time_ms(10, 50);
                assert!(ms >= 0.0);
            }
            Err(_) => {
                // Model not in cache — skip; CI runner may or may not have it
            }
        }
    }

    /// When model is cached: embed an empty batch must return immediately with
    /// no embeddings (the `texts.is_empty()` fast-path in embed_batch_internal).
    #[tokio::test]
    async fn test_engine_embed_empty_batch_when_cached() {
        let config = ModelConfig::new(EmbeddingModel::MiniLM);
        if let Ok(engine) = EmbeddingEngine::new(config).await {
            let result = engine.embed_raw(&[]).await;
            assert!(result.is_ok());
            assert!(result.unwrap().is_empty());
        }
    }

    // ── Session pool (DAK-5547) ──────────────────────────────────────────────

    #[test]
    fn test_session_pool_default_is_4() {
        // Default pool size is 4 (DAK-5746: pool=1 regressed LME ingest ~4×; OOM causes fixed).
        // DAKERA_ONNX_POOL_SIZE env var still allows override.
        let config = ModelConfig::default();
        let expected = std::env::var("DAKERA_ONNX_POOL_SIZE")
            .ok()
            .and_then(|v| v.parse::<usize>().ok())
            .filter(|&n| n >= 1)
            .unwrap_or(4);
        assert_eq!(config.session_pool_size, expected);
    }

    #[test]
    fn test_session_pool_size_builder_roundtrip() {
        let builder = EmbeddingEngineBuilder::new().session_pool_size(8);
        assert_eq!(builder.config.session_pool_size, 8);
    }

    #[test]
    fn test_session_pool_size_min_enforced() {
        let builder = EmbeddingEngineBuilder::new().session_pool_size(0);
        assert_eq!(
            builder.config.session_pool_size, 1,
            "pool size 0 must clamp to 1"
        );
    }

    #[test]
    fn test_model_config_with_session_pool_size() {
        let cfg = ModelConfig::new(EmbeddingModel::MiniLM).with_session_pool_size(2);
        assert_eq!(cfg.session_pool_size, 2);
    }

    /// When model is cached: verify the session pool has the expected size.
    #[tokio::test]
    async fn test_engine_pool_size_matches_config_when_cached() {
        let config = ModelConfig::new(EmbeddingModel::MiniLM).with_session_pool_size(2);
        if let Ok(engine) = EmbeddingEngine::new(config).await {
            assert_eq!(
                engine.pool_size(),
                2,
                "engine should hold exactly 2 sessions"
            );
        }
    }

    // ── next_session round-robin ──────────────────────────────────────────────

    /// Round-robin counter distributes batches across pool slots without panic.
    #[test]
    fn test_round_robin_index_stays_in_bounds() {
        let pool_len = 4_usize;
        let counter = AtomicUsize::new(0);
        for expected_idx in 0..100_usize {
            let start = counter.fetch_add(1, Ordering::Relaxed);
            let slot = start % pool_len;
            assert!(slot < pool_len);
            assert_eq!(slot, expected_idx % pool_len);
        }
    }

    /// Pool size 1 degrades to single-session behavior without panicking.
    #[test]
    fn test_round_robin_pool_size_one() {
        let pool_len = 1_usize;
        let counter = AtomicUsize::new(0);
        for _ in 0..50 {
            let start = counter.fetch_add(1, Ordering::Relaxed);
            assert_eq!(start % pool_len, 0);
        }
    }

    /// When model is cached: empty batch short-circuits before touching the pool.
    #[tokio::test]
    async fn test_embed_empty_does_not_advance_pool_counter() {
        let config = ModelConfig::new(EmbeddingModel::MiniLM).with_session_pool_size(2);
        if let Ok(engine) = EmbeddingEngine::new(config).await {
            let result = engine.embed_raw(&[]).await;
            assert!(result.is_ok());
            assert!(result.unwrap().is_empty());
            // Empty batch returns before fetch_add — counter stays at 0.
            assert_eq!(engine.next_session.load(Ordering::Relaxed), 0);
        }
    }
}