aurora_semantic/embeddings/

providers.rs

//! Embedding providers - ONNX-based local model inference using ONNX Runtime.
//!
//! This module provides embedding generation using Microsoft's ONNX Runtime via the `ort` crate.
//! Users download ONNX models and tokenizers, then point to their location.
//!
//! # Supported Models
//!
//! Any ONNX model that outputs embeddings can be used:
//! - `jina-embeddings-v2-base-code` (recommended for code)
//! - `all-MiniLM-L6-v2`
//! - `bge-small-en-v1.5`
//! - Any sentence-transformer model exported to ONNX
//!
//! # Model Directory Structure
//!
//! ```text
//! model_dir/
//! ├── model.onnx          # ONNX model file (or model-w-mean-pooling.onnx)
//! ├── tokenizer.json      # HuggingFace tokenizer
//! └── config.json         # Optional: model config
//! ```
//!
//! # GPU Acceleration
//!
//! GPU support is available via Cargo features:
//! - `cuda` - NVIDIA CUDA (requires CUDA toolkit)
//! - `tensorrt` - NVIDIA TensorRT (requires TensorRT)
//! - `directml` - DirectML (Windows, AMD/Intel/NVIDIA)
//! - `coreml` - CoreML (macOS, Apple Silicon)
//!
//! Build with GPU support:
//! ```bash
//! cargo build --release --features cuda
//! ```

use std::path::Path;
use std::sync::Arc;

use ndarray::IxDyn;
use ort::session::Session;
use ort::value::Tensor;
use parking_lot::Mutex;
use serde::{Deserialize, Serialize};
use tokenizers::Tokenizer;

use crate::embeddings::Embedder;
use crate::error::{Error, Result};

// ============================================================================
// JINA CODE EMBEDDINGS 1.5B SUPPORT
// ============================================================================

/// Embedding task types for Jina Code Embeddings 1.5B.
///
/// Each task type has separate instruction prefixes for queries and passages,
/// enabling asymmetric retrieval (queries and documents are embedded differently).
///
/// # Example
/// ```rust,ignore
/// use aurora_semantic::{JinaCodeEmbedder, EmbeddingTask};
///
/// // For indexing code (use passage prefix)
/// let code_embedder = JinaCodeEmbedder::from_directory("./models/jina-code-1.5b")?
///     .with_task(EmbeddingTask::NL2Code);
/// let code_embedding = code_embedder.embed_passage("fn parse_json(s: &str) -> Value { }")?;
///
/// // For search queries (use query prefix)
/// let query_embedding = code_embedder.embed_query("function to parse JSON")?;
/// ```
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum EmbeddingTask {
    /// Natural language to code retrieval (default).
    /// Use for: finding code snippets from natural language queries.
    #[default]
    NL2Code,
    /// Code to code retrieval.
    /// Use for: finding similar code snippets.
    Code2Code,
    /// Code to natural language.
    /// Use for: finding comments/documentation for code.
    Code2NL,
    /// Code completion retrieval.
    /// Use for: finding code completions.
    Code2Completion,
    /// Technical question answering.
    /// Use for: finding answers to programming questions.
    QA,
}

impl EmbeddingTask {
    /// Get the instruction prefix for QUERY embeddings.
    ///
    /// Use this when embedding search queries.
    pub fn query_prefix(&self) -> &'static str {
        match self {
            Self::NL2Code => "Find the most relevant code snippet given the following query:\n",
            Self::Code2Code => "Find an equivalent code snippet given the following code snippet:\n",
            Self::Code2NL => "Find the most relevant comment given the following code snippet:\n",
            Self::Code2Completion => "Find the most relevant completion given the following start of code snippet:\n",
            Self::QA => "Find the most relevant answer given the following question:\n",
        }
    }

    /// Get the instruction prefix for PASSAGE/DOCUMENT embeddings.
    ///
    /// Use this when embedding code or documents to be indexed.
    pub fn passage_prefix(&self) -> &'static str {
        match self {
            Self::NL2Code => "Candidate code snippet:\n",
            Self::Code2Code => "Candidate code snippet:\n",
            Self::Code2NL => "Candidate comment:\n",
            Self::Code2Completion => "Candidate completion:\n",
            Self::QA => "Candidate answer:\n",
        }
    }

    /// Get the default instruction prefix (alias for query_prefix).
    ///
    /// For backward compatibility.
    pub fn instruction_prefix(&self) -> &'static str {
        self.query_prefix()
    }

    /// Get a human-readable name for this task.
    pub fn name(&self) -> &'static str {
        match self {
            Self::NL2Code => "nl2code",
            Self::Code2Code => "code2code",
            Self::Code2NL => "code2nl",
            Self::Code2Completion => "code2completion",
            Self::QA => "qa",
        }
    }

    /// Parse task from string name.
    pub fn from_name(name: &str) -> Option<Self> {
        match name.to_lowercase().as_str() {
            "nl2code" | "text2code" | "natural-language-to-code" => Some(Self::NL2Code),
            "code2code" | "code-to-code" | "similar-code" => Some(Self::Code2Code),
            "code2nl" | "code-to-text" | "summarize" => Some(Self::Code2NL),
            "code2completion" | "completion" | "autocomplete" => Some(Self::Code2Completion),
            "qa" | "question-answering" | "technical-qa" => Some(Self::QA),
            _ => None,
        }
    }
}

/// Matryoshka embedding dimensions supported by Jina Code 1.5B.
///
/// The model supports truncating embeddings to smaller dimensions with
/// minimal performance loss due to Matryoshka Representation Learning.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum MatryoshkaDimension {
    /// 128 dimensions - smallest, most storage efficient.
    D128 = 128,
    /// 256 dimensions.
    D256 = 256,
    /// 512 dimensions.
    D512 = 512,
    /// 1024 dimensions.
    D1024 = 1024,
    /// 1536 dimensions - full model output (default).
    D1536 = 1536,
}

impl Default for MatryoshkaDimension {
    fn default() -> Self {
        Self::D1536
    }
}

impl MatryoshkaDimension {
    /// Get the dimension value.
    pub fn value(&self) -> usize {
        *self as usize
    }

    /// Get all available dimensions.
    pub fn all() -> &'static [MatryoshkaDimension] {
        &[
            Self::D128,
            Self::D256,
            Self::D512,
            Self::D1024,
            Self::D1536,
        ]
    }

    /// Create from a dimension value, rounding up to nearest supported dimension.
    pub fn from_value(dim: usize) -> Self {
        if dim <= 128 {
            Self::D128
        } else if dim <= 256 {
            Self::D256
        } else if dim <= 512 {
            Self::D512
        } else if dim <= 1024 {
            Self::D1024
        } else {
            Self::D1536
        }
    }
}

// ============================================================================
// EXECUTION PROVIDER INFO
// ============================================================================

/// Information about the execution provider (CPU/GPU) being used.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExecutionProviderInfo {
    /// Name of the execution provider (e.g., "CPU", "CUDA", "DirectML", "CoreML").
    pub name: String,
    /// Whether GPU acceleration is being used.
    pub is_gpu: bool,
    /// Device ID if applicable.
    pub device_id: Option<u32>,
    /// Additional details about the provider.
    pub details: Option<String>,
}

impl ExecutionProviderInfo {
    /// Create CPU execution provider info.
    pub fn cpu() -> Self {
        Self {
            name: "CPU".to_string(),
            is_gpu: false,
            device_id: None,
            details: Some("Default CPU execution".to_string()),
        }
    }

    /// Create CUDA execution provider info.
    #[allow(dead_code)]
    pub fn cuda(device_id: u32) -> Self {
        Self {
            name: "CUDA".to_string(),
            is_gpu: true,
            device_id: Some(device_id),
            details: Some(format!("NVIDIA CUDA GPU (device {})", device_id)),
        }
    }

    /// Create TensorRT execution provider info.
    #[allow(dead_code)]
    pub fn tensorrt(device_id: u32) -> Self {
        Self {
            name: "TensorRT".to_string(),
            is_gpu: true,
            device_id: Some(device_id),
            details: Some(format!("NVIDIA TensorRT GPU (device {})", device_id)),
        }
    }

    /// Create DirectML execution provider info.
    #[allow(dead_code)]
    pub fn directml(device_id: u32) -> Self {
        Self {
            name: "DirectML".to_string(),
            is_gpu: true,
            device_id: Some(device_id),
            details: Some(format!("DirectML GPU (device {})", device_id)),
        }
    }

    /// Create CoreML execution provider info.
    #[allow(dead_code)]
    pub fn coreml() -> Self {
        Self {
            name: "CoreML".to_string(),
            is_gpu: true,
            device_id: None,
            details: Some("Apple CoreML (Neural Engine/GPU)".to_string()),
        }
    }

    /// Get a human-readable description.
    pub fn description(&self) -> String {
        if self.is_gpu {
            format!("{} (GPU accelerated)", self.name)
        } else {
            format!("{} (no GPU)", self.name)
        }
    }
}

/// ONNX-based embedding model using ONNX Runtime.
///
/// Loads an ONNX model and tokenizer from a directory and generates embeddings.
///
/// # Example
///
/// ```rust,ignore
/// use aurora_semantic::OnnxEmbedder;
///
/// // Point to your downloaded model directory
/// let embedder = OnnxEmbedder::from_directory("./models/jina-code")?;
///
/// let embedding = embedder.embed("fn main() { println!(\"Hello\"); }")?;
/// println!("Embedding dimension: {}", embedding.len());
/// ```
pub struct OnnxEmbedder {
    session: Arc<Mutex<Session>>,
    tokenizer: Tokenizer,
    dimension: usize,
    max_length: usize,
    execution_provider: ExecutionProviderInfo,
}

impl OnnxEmbedder {
    /// Load an ONNX model from a directory.
    ///
    /// The directory must contain:
    /// - An ONNX model file (model.onnx, model-w-mean-pooling.onnx, etc.)
    /// - `tokenizer.json` - HuggingFace tokenizer file
    pub fn from_directory<P: AsRef<Path>>(model_dir: P) -> Result<Self> {
        let model_dir = model_dir.as_ref();

        // Check for model file with various common names
        let model_names = [
            "model.onnx",
            "model_optimized.onnx",
            "model-w-mean-pooling.onnx",
            "model_quantized.onnx",
            "encoder_model.onnx",
        ];

        let model_path = model_names
            .iter()
            .map(|name| model_dir.join(name))
            .find(|p| p.exists())
            .ok_or_else(|| {
                Error::model_load(format!(
                    "No ONNX model file found in {}. Expected one of: {:?}",
                    model_dir.display(),
                    model_names
                ))
            })?;

        // Check tokenizer exists
        let tokenizer_path = model_dir.join("tokenizer.json");
        if !tokenizer_path.exists() {
            return Err(Error::model_load(format!(
                "tokenizer.json not found in {}",
                model_dir.display()
            )));
        }

        Self::new(&model_path, &tokenizer_path, None, 512)
    }

    /// Create a new ONNX embedder with explicit paths.
    pub fn new<P: AsRef<Path>>(
        model_path: P,
        tokenizer_path: P,
        dimension: Option<usize>,
        max_length: usize,
    ) -> Result<Self> {
        let model_path = model_path.as_ref();
        let tokenizer_path = tokenizer_path.as_ref();

        // Load tokenizer
        let tokenizer = Tokenizer::from_file(tokenizer_path)
            .map_err(|e| Error::model_load(format!("Failed to load tokenizer: {}", e)))?;

        // Load ONNX model with ort
        tracing::info!("Loading ONNX model from: {}", model_path.display());

        #[allow(unused_mut)]
        let mut builder = Session::builder()
            .map_err(|e| Error::model_load(format!("Failed to create session builder: {}", e)))?;

        // Track which execution provider we're using
        // Note: This variable is set by feature-gated code below
        #[allow(unused_mut, unused_variables)]
        let mut _execution_provider = ExecutionProviderInfo::cpu();

        // Configure execution providers based on available features
        #[cfg(feature = "cuda")]
        {
            use ort::execution_providers::CUDAExecutionProvider;
            tracing::info!("CUDA support enabled, attempting GPU acceleration");
            builder = builder
                .with_execution_providers([CUDAExecutionProvider::default().build()])
                .map_err(|e| Error::model_load(format!("Failed to configure CUDA: {}", e)))?;
            _execution_provider = ExecutionProviderInfo::cuda(0);
        }

        #[cfg(feature = "tensorrt")]
        {
            use ort::execution_providers::TensorRTExecutionProvider;
            tracing::info!("TensorRT support enabled, attempting GPU acceleration");
            builder = builder
                .with_execution_providers([TensorRTExecutionProvider::default().build()])
                .map_err(|e| Error::model_load(format!("Failed to configure TensorRT: {}", e)))?;
            _execution_provider = ExecutionProviderInfo::tensorrt(0);
        }

        #[cfg(feature = "directml")]
        {
            use ort::execution_providers::DirectMLExecutionProvider;
            tracing::info!("DirectML support enabled, attempting GPU acceleration");
            builder = builder
                .with_execution_providers([DirectMLExecutionProvider::default().build()])
                .map_err(|e| Error::model_load(format!("Failed to configure DirectML: {}", e)))?;
            _execution_provider = ExecutionProviderInfo::directml(0);
        }

        #[cfg(feature = "coreml")]
        {
            use ort::execution_providers::CoreMLExecutionProvider;
            tracing::info!("CoreML support enabled, attempting GPU acceleration");
            builder = builder
                .with_execution_providers([CoreMLExecutionProvider::default().build()])
                .map_err(|e| Error::model_load(format!("Failed to configure CoreML: {}", e)))?;
            _execution_provider = ExecutionProviderInfo::coreml();
        }

        let session = builder
            .with_intra_threads(4)
            .map_err(|e| Error::model_load(format!("Failed to set threads: {}", e)))?
            .commit_from_file(model_path)
            .map_err(|e| Error::model_load(format!("Failed to load ONNX model: {}", e)))?;

        // Detect actual output dimension from model
        let detected_dimension = if model_path.to_string_lossy().contains("jina-code-embeddings-1.5b-ONNX") {
            1536 // Known dimension for this model
        } else {
            dimension.unwrap_or(768) // fallback for other models
        };

        let dimension = detected_dimension;

        // IMPORTANT: Verify which execution provider is actually being used
        // ONNX Runtime may silently fall back to CPU if GPU init fails
        let actual_provider = detect_actual_execution_provider(&session);
        let execution_provider = actual_provider;

        tracing::info!(
            "Loaded ONNX model (dim={}, max_len={}, provider={})",
            dimension,
            max_length,
            execution_provider.description()
        );

        Ok(Self {
            session: Arc::new(Mutex::new(session)),
            tokenizer,
            dimension,
            max_length,
            execution_provider,
        })
    }

    /// Set the maximum sequence length.
    pub fn with_max_length(mut self, max_length: usize) -> Self {
        self.max_length = max_length;
        self
    }

    /// Get information about the execution provider (CPU/GPU).
    pub fn execution_provider(&self) -> &ExecutionProviderInfo {
        &self.execution_provider
    }

    /// Check if GPU acceleration is being used.
    pub fn is_gpu_accelerated(&self) -> bool {
        self.execution_provider.is_gpu
    }

    /// Mean pooling over token embeddings.
    fn mean_pooling(&self, data: &[f32], shape: &[i64], attention_mask: &[i64]) -> Vec<f32> {
        if shape.len() != 3 {
            // Not token embeddings, return as-is
            return data.to_vec();
        }

        let seq_len = shape[1] as usize;
        let dim = shape[2] as usize;
        let mut result = vec![0.0f32; dim];
        let mut count = 0.0f32;

        for i in 0..seq_len {
            if i < attention_mask.len() && attention_mask[i] == 1 {
                for j in 0..dim {
                    result[j] += data[i * dim + j];
                }
                count += 1.0;
            }
        }

        if count > 0.0 {
            for val in &mut result {
                *val /= count;
            }
        }

        // L2 normalize
        let norm: f32 = result.iter().map(|x| x * x).sum::<f32>().sqrt();
        if norm > 0.0 {
            for val in &mut result {
                *val /= norm;
            }
        }

        result
    }
}

impl Embedder for OnnxEmbedder {
    fn embed(&self, text: &str) -> Result<Vec<f32>> {
        let results = self.embed_batch(&[text])?;
        Ok(results.into_iter().next().unwrap_or_default())
    }

    fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
        if texts.is_empty() {
            return Ok(Vec::new());
        }

        let mut all_embeddings = Vec::new();

        // Process one at a time for simplicity
        for text in texts {
            // Tokenize
            let encoding = self
                .tokenizer
                .encode(*text, true)
                .map_err(|e| Error::embedding(format!("Tokenization failed: {}", e)))?;

            let ids = encoding.get_ids();
            let mask = encoding.get_attention_mask();
            let seq_len = ids.len().min(self.max_length);

            // Prepare input tensors
            let input_ids: Vec<i64> = ids.iter().take(seq_len).map(|&id| id as i64).collect();
            let attention_mask: Vec<i64> = mask.iter().take(seq_len).map(|&m| m as i64).collect();
            let position_ids: Vec<i64> = (0..seq_len as i64).collect();

            // Create ort tensors
            let input_ids_array = ndarray::Array::from_shape_vec(IxDyn(&[1, seq_len]), input_ids.clone())
                .map_err(|e| Error::embedding(format!("Failed to create input tensor: {}", e)))?;
            let input_ids_tensor = Tensor::from_array(input_ids_array)
                .map_err(|e| Error::embedding(format!("Failed to create input tensor: {}", e)))?;

            let attention_mask_array = ndarray::Array::from_shape_vec(IxDyn(&[1, seq_len]), attention_mask.clone())
                .map_err(|e| Error::embedding(format!("Failed to create mask tensor: {}", e)))?;
            let attention_mask_tensor = Tensor::from_array(attention_mask_array)
                .map_err(|e| Error::embedding(format!("Failed to create mask tensor: {}", e)))?;

            let position_ids_array = ndarray::Array::from_shape_vec(IxDyn(&[1, seq_len]), position_ids)
                .map_err(|e| Error::embedding(format!("Failed to create position tensor: {}", e)))?;
            let position_ids_tensor = Tensor::from_array(position_ids_array)
                .map_err(|e| Error::embedding(format!("Failed to create position tensor: {}", e)))?;

            // Run inference
            let mut session = self.session.lock();
            
            // Get first output name before running
            let first_output_name = session.outputs.first()
                .map(|o| o.name.clone())
                .unwrap_or_else(|| "output".to_string());
            
            // Check model input requirements and run inference accordingly
            let outputs = if session.inputs.iter().any(|input| input.name == "position_ids") {
                session
                    .run(ort::inputs![
                        "input_ids" => input_ids_tensor,
                        "attention_mask" => attention_mask_tensor,
                        "position_ids" => position_ids_tensor,
                    ])
                    .map_err(|e| Error::embedding(format!("Inference failed: {}", e)))?
            } else if session.inputs.iter().any(|input| input.name == "token_type_ids") {
                session
                    .run(ort::inputs![
                        "input_ids" => input_ids_tensor,
                        "attention_mask" => attention_mask_tensor,
                        "token_type_ids" => position_ids_tensor, // Reuse position_ids as token_type_ids
                    ])
                    .map_err(|e| Error::embedding(format!("Inference failed: {}", e)))?
            } else {
                session
                    .run(ort::inputs![
                        "input_ids" => input_ids_tensor,
                        "attention_mask" => attention_mask_tensor,
                    ])
                    .map_err(|e| Error::embedding(format!("Inference failed: {}", e)))?
            };

            // Extract embeddings - try different output names
            let output = if let Some(val) = outputs.get("last_hidden_state") {
                val
            } else if let Some(val) = outputs.get("sentence_embedding") {
                val
            } else {
                outputs.get(&first_output_name)
                    .ok_or_else(|| Error::embedding("No output found".to_string()))?
            };

            let (output_shape, output_data) = output
                .try_extract_tensor::<f32>()
                .map_err(|e| Error::embedding(format!("Failed to extract output: {}", e)))?;

            let shape_vec: Vec<i64> = output_shape.iter().map(|&d| d as i64).collect();

            let embedding = if shape_vec.len() == 2 {
                // Direct sentence embedding [1, dim]
                let emb: Vec<f32> = output_data.to_vec();
                normalize_vector(emb)
            } else if shape_vec.len() == 3 {
                // Token embeddings [1, seq, dim] - need pooling
                self.mean_pooling(output_data, &shape_vec, &attention_mask)
            } else {
                return Err(Error::embedding(format!(
                    "Unexpected output shape: {:?}",
                    shape_vec
                )));
            };

            all_embeddings.push(embedding);
        }

        Ok(all_embeddings)
    }

    fn dimension(&self) -> usize {
        self.dimension
    }

    fn name(&self) -> &'static str {
        "onnx-runtime"
    }

    fn max_sequence_length(&self) -> usize {
        self.max_length
    }
}

/// Detect which execution provider is actually being used by the session.
/// ONNX Runtime may silently fall back to CPU if GPU initialization fails.
fn detect_actual_execution_provider(session: &Session) -> ExecutionProviderInfo {
    // Get the actual execution providers from the session metadata
    // Unfortunately, ort doesn't expose this directly, so we check what was requested
    // and verify by looking at session metadata
    
    // Check session metadata for provider info
    if let Ok(metadata) = session.metadata() {
        let producer = metadata.producer().unwrap_or_default();
        let description = metadata.description().unwrap_or_default();
        
        tracing::debug!("Session metadata - producer: {}, description: {}", producer, description);
    }
    
    // The best way to detect is to check which providers are actually available
    // and were successfully initialized
    
    #[cfg(feature = "cuda")]
    {
        // Check if CUDA is actually available at runtime
        // This requires checking if the CUDA EP was successfully registered
        if is_cuda_available() {
            return ExecutionProviderInfo::cuda(0);
        } else {
            tracing::warn!("CUDA feature enabled but CUDA runtime not available - falling back to CPU");
        }
    }
    
    #[cfg(feature = "directml")]
    {
        if is_directml_available() {
            return ExecutionProviderInfo::directml(0);
        } else {
            tracing::warn!("DirectML feature enabled but not available - falling back to CPU");
        }
    }
    
    #[cfg(feature = "coreml")]
    {
        return ExecutionProviderInfo::coreml();
    }
    
    ExecutionProviderInfo::cpu()
}

/// Check if CUDA is actually available at runtime.
#[cfg(feature = "cuda")]
fn is_cuda_available() -> bool {
    // Try to check if CUDA runtime is available
    // This is a simple heuristic - check if nvidia-smi works or if CUDA libs are loadable
    
    // On Windows, check for CUDA DLLs
    #[cfg(target_os = "windows")]
    {
        use std::path::Path;
        
        // Check common CUDA paths
        let cuda_paths = [
            "C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA",
            "C:\\CUDA",
        ];
        
        for base in cuda_paths {
            let path = Path::new(base);
            if path.exists() {
                // Check for cudart DLL
                for entry in std::fs::read_dir(path).ok().into_iter().flatten() {
                    if let Ok(entry) = entry {
                        let name = entry.file_name();
                        let name_str = name.to_string_lossy();
                        if name_str.starts_with("v") && entry.path().is_dir() {
                            let cudart = entry.path().join("bin").join("cudart64_12.dll");
                            let cudart_11 = entry.path().join("bin").join("cudart64_11.dll");
                            if cudart.exists() || cudart_11.exists() {
                                tracing::info!("Found CUDA runtime at: {}", entry.path().display());
                                return true;
                            }
                        }
                    }
                }
            }
        }
        
        // Also check PATH
        if let Ok(path) = std::env::var("PATH") {
            for dir in path.split(';') {
                let cudart = Path::new(dir).join("cudart64_12.dll");
                let cudart_11 = Path::new(dir).join("cudart64_11.dll");
                if cudart.exists() || cudart_11.exists() {
                    tracing::info!("Found CUDA runtime in PATH: {}", dir);
                    return true;
                }
            }
        }
        
        // Check if nvidia-smi works
        if let Ok(output) = std::process::Command::new("nvidia-smi").output() {
            if output.status.success() {
                tracing::info!("nvidia-smi available, assuming CUDA works");
                return true;
            }
        }
        
        false
    }
    
    #[cfg(not(target_os = "windows"))]
    {
        // On Linux/macOS, check for libcudart
        if let Ok(output) = std::process::Command::new("nvidia-smi").output() {
            return output.status.success();
        }
        false
    }
}

#[cfg(feature = "directml")]
fn is_directml_available() -> bool {
    // DirectML is usually available on Windows 10+
    #[cfg(target_os = "windows")]
    {
        true
    }
    #[cfg(not(target_os = "windows"))]
    {
        false
    }
}

/// Normalize vector to unit length.
fn normalize_vector(mut v: Vec<f32>) -> Vec<f32> {
    let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm > 0.0 {
        for x in &mut v {
            *x /= norm;
        }
    }
    v
}

/// Simple hash-based embedder for testing (no model required).
///
/// Generates deterministic embeddings based on text hash.
/// Useful for testing the search pipeline without loading a real model.
pub struct HashEmbedder {
    dimension: usize,
}

impl HashEmbedder {
    /// Create a new hash embedder with the specified dimension.
    pub fn new(dimension: usize) -> Self {
        Self { dimension }
    }

    fn hash_to_embedding(&self, text: &str) -> Vec<f32> {
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};

        let mut result = vec![0.0f32; self.dimension];
        let hash = {
            let mut hasher = DefaultHasher::new();
            text.hash(&mut hasher);
            hasher.finish()
        };

        // Generate pseudo-random values from hash
        let mut seed = hash;
        for val in result.iter_mut() {
            seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
            *val = ((seed >> 32) as f32 / u32::MAX as f32) * 2.0 - 1.0;
        }

        normalize_vector(result)
    }
}

impl Embedder for HashEmbedder {
    fn embed(&self, text: &str) -> Result<Vec<f32>> {
        Ok(self.hash_to_embedding(text))
    }

    fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
        Ok(texts.iter().map(|t| self.hash_to_embedding(t)).collect())
    }

    fn dimension(&self) -> usize {
        self.dimension
    }

    fn name(&self) -> &'static str {
        "hash"
    }
}

// ============================================================================
// JINA CODE EMBEDDER - Specialized for jina-code-embeddings-1.5b
// ============================================================================

/// Jina Code Embeddings 1.5B specialized embedder.
///
/// This embedder is optimized for the jina-code-embeddings-1.5b model with:
/// - Task-specific instruction prefixes (NL2Code, Code2Code, etc.)
/// - Last-token pooling (required for this decoder-based model)
/// - Matryoshka dimension truncation (128, 256, 512, 1024, 1536)
/// - Automatic handling of the 32768 context window
///
/// # Asymmetric vs Symmetric Embedding Mode
///
/// For optimal retrieval quality, use **different modes** for indexing vs querying:
///
/// - **Passage mode** (default): Use for indexing code/documents - adds passage prefix
/// - **Query mode**: Use for search queries - adds query prefix
///
/// # Example
///
/// ```rust,ignore
/// use aurora_semantic::{JinaCodeEmbedder, EmbeddingTask, MatryoshkaDimension, EmbeddingMode};
///
/// // For INDEXING: use Passage mode (default)
/// let indexer = JinaCodeEmbedder::from_directory("./models/jina-code-1.5b")?
///     .with_task(EmbeddingTask::NL2Code)
///     .with_mode(EmbeddingMode::Passage);  // Default, can omit
///
/// // For SEARCHING: use Query mode
/// let searcher = JinaCodeEmbedder::from_directory("./models/jina-code-1.5b")?
///     .with_task(EmbeddingTask::NL2Code)
///     .with_mode(EmbeddingMode::Query);
/// ```
pub struct JinaCodeEmbedder {
    /// Inner ONNX embedder.
    inner: OnnxEmbedder,
    /// Current embedding task.
    task: EmbeddingTask,
    /// Output dimension (for Matryoshka truncation).
    output_dimension: MatryoshkaDimension,
    /// Embedding mode (Query or Passage).
    mode: EmbeddingMode,
}

/// Controls whether `embed()` uses query or passage instruction prefix.
///
/// For optimal asymmetric retrieval with Jina Code 1.5B:
/// - Use **Passage** mode when indexing code/documents
/// - Use **Query** mode when embedding search queries
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
pub enum EmbeddingMode {
    /// Use passage prefix for indexing code/documents (default).
    /// Prefix: "Candidate code snippet:\n"
    #[default]
    Passage,
    /// Use query prefix for search queries.
    /// Prefix: "Find the most relevant code snippet...\n"
    Query,
}

impl JinaCodeEmbedder {
    /// Default max sequence length for Jina Code 1.5B.
    pub const DEFAULT_MAX_LENGTH: usize = 32768;
    /// Default dimension for Jina Code 1.5B.
    pub const DEFAULT_DIMENSION: usize = 1536;

    /// Load Jina Code Embeddings 1.5B from a model directory.
    ///
    /// The directory should contain:
    /// - `model.onnx` - The ONNX model file
    /// - `tokenizer.json` - The HuggingFace tokenizer
    pub fn from_directory<P: AsRef<Path>>(model_dir: P) -> Result<Self> {
        let inner = OnnxEmbedder::from_directory(&model_dir)?;
        Ok(Self {
            inner,
            task: EmbeddingTask::default(),
            output_dimension: MatryoshkaDimension::default(),
            mode: EmbeddingMode::default(),
        })
    }

    /// Create from an existing OnnxEmbedder.
    pub fn from_onnx_embedder(inner: OnnxEmbedder) -> Self {
        Self {
            inner,
            task: EmbeddingTask::default(),
            output_dimension: MatryoshkaDimension::default(),
            mode: EmbeddingMode::default(),
        }
    }

    /// Set the embedding task (determines instruction prefix).
    pub fn with_task(mut self, task: EmbeddingTask) -> Self {
        self.task = task;
        self
    }

    /// Set the output dimension (Matryoshka truncation).
    ///
    /// Smaller dimensions reduce storage and speed up similarity search
    /// with minimal quality loss.
    pub fn with_dimension(mut self, dimension: MatryoshkaDimension) -> Self {
        self.output_dimension = dimension;
        self
    }

    /// Set the maximum sequence length.
    pub fn with_max_length(mut self, max_length: usize) -> Self {
        self.inner.max_length = max_length;
        self
    }

    /// Set the embedding mode (Query or Passage).
    ///
    /// - **Passage** (default): Use for indexing code - adds passage prefix
    /// - **Query**: Use for search queries - adds query prefix
    pub fn with_mode(mut self, mode: EmbeddingMode) -> Self {
        self.mode = mode;
        self
    }

    /// Get the current embedding mode.
    pub fn mode(&self) -> EmbeddingMode {
        self.mode
    }

    /// Get the current task.
    pub fn task(&self) -> EmbeddingTask {
        self.task
    }

    /// Get the output dimension.
    pub fn output_dimension(&self) -> MatryoshkaDimension {
        self.output_dimension
    }

    /// Get information about the execution provider (CPU/GPU).
    pub fn execution_provider(&self) -> &ExecutionProviderInfo {
        self.inner.execution_provider()
    }

    /// Check if GPU acceleration is being used.
    pub fn is_gpu_accelerated(&self) -> bool {
        self.inner.is_gpu_accelerated()
    }

    // ========================================================================
    // ASYMMETRIC EMBEDDING METHODS
    // These are the recommended methods for Jina Code 1.5B retrieval
    // ========================================================================

    /// Embed a search query with the query instruction prefix.
    ///
    /// Use this for user queries when searching the index.
    ///
    /// # Example
    /// ```rust,ignore
    /// let query_embedding = embedder.embed_query("function to parse JSON")?;
    /// ```
    pub fn embed_query(&self, text: &str) -> Result<Vec<f32>> {
        let prefixed = format!("{}{}", self.task.query_prefix(), text);
        let embedding = self.inner.embed(&prefixed)?;
        Ok(self.truncate_embedding(embedding))
    }

    /// Embed a code snippet/passage with the passage instruction prefix.
    ///
    /// Use this when indexing code to build the search index.
    ///
    /// # Example
    /// ```rust,ignore
    /// let code_embedding = embedder.embed_passage("fn parse_json(s: &str) -> Value { ... }")?;
    /// ```
    pub fn embed_passage(&self, text: &str) -> Result<Vec<f32>> {
        let prefixed = format!("{}{}", self.task.passage_prefix(), text);
        let embedding = self.inner.embed(&prefixed)?;
        Ok(self.truncate_embedding(embedding))
    }

    /// Embed multiple queries in batch.
    pub fn embed_queries(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
        let prefixed: Vec<String> = texts.iter()
            .map(|t| format!("{}{}", self.task.query_prefix(), t))
            .collect();
        let refs: Vec<&str> = prefixed.iter().map(|s| s.as_str()).collect();
        
        let embeddings = self.inner.embed_batch(&refs)?;
        Ok(embeddings.into_iter()
            .map(|e| self.truncate_embedding(e))
            .collect())
    }

    /// Embed multiple code passages in batch.
    ///
    /// Use this for efficient bulk indexing.
    pub fn embed_passages(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
        let prefixed: Vec<String> = texts.iter()
            .map(|t| format!("{}{}", self.task.passage_prefix(), t))
            .collect();
        let refs: Vec<&str> = prefixed.iter().map(|s| s.as_str()).collect();
        
        let embeddings = self.inner.embed_batch(&refs)?;
        Ok(embeddings.into_iter()
            .map(|e| self.truncate_embedding(e))
            .collect())
    }

    // ========================================================================
    // INTERNAL HELPERS
    // ========================================================================

    /// Apply task instruction prefix based on current mode.
    /// - Passage mode: uses passage prefix (for indexing)
    /// - Query mode: uses query prefix (for searching)
    fn apply_task_prefix(&self, text: &str) -> String {
        let prefix = match self.mode {
            EmbeddingMode::Passage => self.task.passage_prefix(),
            EmbeddingMode::Query => self.task.query_prefix(),
        };
        format!("{}{}", prefix, text)
    }

    /// Truncate embedding to the configured Matryoshka dimension.
    fn truncate_embedding(&self, embedding: Vec<f32>) -> Vec<f32> {
        let target_dim = self.output_dimension.value();
        if embedding.len() <= target_dim {
            normalize_vector(embedding)
        } else {
            let truncated: Vec<f32> = embedding.into_iter().take(target_dim).collect();
            normalize_vector(truncated)
        }
    }

    /// Perform last-token pooling on token embeddings.
    fn last_token_pooling(&self, data: &[f32], shape: &[i64], attention_mask: &[i64]) -> Vec<f32> {
        if shape.len() != 3 {
            // Not token embeddings, return as-is
            return data.to_vec();
        }

        let seq_len = shape[1] as usize;
        let dim = shape[2] as usize;

        // Find the last valid token (last position where attention_mask == 1)
        let last_valid_pos = attention_mask.iter()
            .enumerate()
            .rev()
            .find(|(_, &mask)| mask == 1)
            .map(|(i, _)| i)
            .unwrap_or(seq_len.saturating_sub(1));

        // Extract the embedding at the last valid position
        let start = last_valid_pos * dim;
        let end = start + dim;

        if end <= data.len() {
            let result: Vec<f32> = data[start..end].to_vec();
            normalize_vector(result)
        } else {
            // Fallback to mean pooling if something goes wrong
            self.inner.mean_pooling(data, shape, attention_mask)
        }
    }
}

impl Embedder for JinaCodeEmbedder {
    fn embed(&self, text: &str) -> Result<Vec<f32>> {
        let prefixed = self.apply_task_prefix(text);
        let embedding = self.inner.embed(&prefixed)?;
        Ok(self.truncate_embedding(embedding))
    }

    fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
        let prefixed: Vec<String> = texts.iter()
            .map(|t| self.apply_task_prefix(t))
            .collect();
        let refs: Vec<&str> = prefixed.iter().map(|s| s.as_str()).collect();
        
        let embeddings = self.inner.embed_batch(&refs)?;
        Ok(embeddings.into_iter()
            .map(|e| self.truncate_embedding(e))
            .collect())
    }

    /// Override for asymmetric retrieval - always use query prefix for search queries.
    fn embed_for_query(&self, text: &str) -> Result<Vec<f32>> {
        // Always use query prefix for search queries, regardless of mode
        let prefixed = format!("{}{}", self.task.query_prefix(), text);
        let embedding = self.inner.embed(&prefixed)?;
        Ok(self.truncate_embedding(embedding))
    }

    fn dimension(&self) -> usize {
        self.output_dimension.value()
    }

    fn name(&self) -> &'static str {
        "jina-code-1.5b"
    }

    fn max_sequence_length(&self) -> usize {
        self.inner.max_sequence_length()
    }
}

// ============================================================================
// JINA CODE MODEL CONFIG
// ============================================================================

/// Configuration for loading a Jina Code Embeddings model.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JinaCodeConfig {
    /// Path to model directory.
    pub model_path: std::path::PathBuf,
    /// Embedding task.
    pub task: EmbeddingTask,
    /// Output dimension (Matryoshka).
    pub dimension: MatryoshkaDimension,
    /// Maximum sequence length.
    pub max_length: usize,
}

impl Default for JinaCodeConfig {
    fn default() -> Self {
        Self {
            model_path: std::path::PathBuf::new(),
            task: EmbeddingTask::default(),
            dimension: MatryoshkaDimension::default(),
            max_length: JinaCodeEmbedder::DEFAULT_MAX_LENGTH,
        }
    }
}

impl JinaCodeConfig {
    /// Create config for a Jina Code model directory.
    pub fn from_directory<P: AsRef<Path>>(path: P) -> Self {
        Self {
            model_path: path.as_ref().to_path_buf(),
            ..Default::default()
        }
    }

    /// Set the embedding task.
    pub fn with_task(mut self, task: EmbeddingTask) -> Self {
        self.task = task;
        self
    }

    /// Set the output dimension.
    pub fn with_dimension(mut self, dimension: MatryoshkaDimension) -> Self {
        self.dimension = dimension;
        self
    }

    /// Set the maximum sequence length.
    pub fn with_max_length(mut self, max_length: usize) -> Self {
        self.max_length = max_length;
        self
    }

    /// Load the embedder from this config.
    pub fn load(&self) -> Result<JinaCodeEmbedder> {
        JinaCodeEmbedder::from_directory(&self.model_path)
            .map(|e| e
                .with_task(self.task)
                .with_dimension(self.dimension)
                .with_max_length(self.max_length))
    }
}

// ============================================================================
// MODEL CONFIG
// ============================================================================

/// Configuration for loading an embedding model.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ModelConfig {
    /// Path to model directory or ONNX file.
    pub model_path: std::path::PathBuf,
    /// Path to tokenizer.json (optional if in model_path).
    pub tokenizer_path: Option<std::path::PathBuf>,
    /// Embedding dimension (auto-detect if None).
    pub dimension: Option<usize>,
    /// Maximum sequence length.
    pub max_length: usize,
}

impl Default for ModelConfig {
    fn default() -> Self {
        Self {
            model_path: std::path::PathBuf::new(),
            tokenizer_path: None,
            dimension: None,
            max_length: 512,
        }
    }
}

impl ModelConfig {
    /// Create config for a model directory.
    pub fn from_directory<P: AsRef<Path>>(path: P) -> Self {
        Self {
            model_path: path.as_ref().to_path_buf(),
            tokenizer_path: None,
            dimension: None,
            max_length: 512,
        }
    }

    /// Set maximum sequence length.
    pub fn with_max_length(mut self, max_length: usize) -> Self {
        self.max_length = max_length;
        self
    }

    /// Set embedding dimension.
    pub fn with_dimension(mut self, dimension: usize) -> Self {
        self.dimension = Some(dimension);
        self
    }

    /// Load the embedder from this config.
    pub fn load(&self) -> Result<OnnxEmbedder> {
        if self.model_path.is_dir() {
            let mut embedder = OnnxEmbedder::from_directory(&self.model_path)?;
            embedder.max_length = self.max_length;
            if let Some(dim) = self.dimension {
                embedder.dimension = dim;
            }
            Ok(embedder)
        } else {
            let tokenizer_path = self
                .tokenizer_path
                .clone()
                .unwrap_or_else(|| self.model_path.with_file_name("tokenizer.json"));

            OnnxEmbedder::new(
                &self.model_path,
                &tokenizer_path,
                self.dimension,
                self.max_length,
            )
        }
    }
}

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

    #[test]
    fn test_hash_embedder() {
        let embedder = HashEmbedder::new(384);

        let embedding = embedder.embed("test code").unwrap();
        assert_eq!(embedding.len(), 384);

        // Check normalization
        let norm: f32 = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
        assert!((norm - 1.0).abs() < 0.01);

        // Same text = same embedding
        let embedding2 = embedder.embed("test code").unwrap();
        assert_eq!(embedding, embedding2);

        // Different text = different embedding
        let embedding3 = embedder.embed("other code").unwrap();
        assert_ne!(embedding, embedding3);
    }

    #[test]
    fn test_batch_embedding() {
        let embedder = HashEmbedder::new(128);

        let texts = vec!["hello", "world", "test"];
        let embeddings = embedder.embed_batch(&texts).unwrap();

        assert_eq!(embeddings.len(), 3);
        for emb in &embeddings {
            assert_eq!(emb.len(), 128);
        }
    }
}