//! Embedding model trait and ONNX integration (feature-flagged).
//!
//! Defines the `EmbeddingModel` trait that all embedding implementations
//! must satisfy. The default `HashEmbedder` wraps the existing hash-based
//! approach. When the `onnx` feature is enabled, an `OnnxEmbedder` can
//! load CodeBERT or other ONNX models.

use crate::search::embedder;

/// Embedding dimension — must match across all implementations.
pub const EMBEDDING_DIM: usize = embedder::EMBEDDING_DIM;

/// Trait for text embedding models.
///
/// Both the hash-based embedder and future neural models implement this
/// trait so they can be swapped transparently in the search pipeline.
pub trait EmbeddingModel: Send + Sync {
    /// Model name for logging/metrics.
    fn name(&self) -> &str;

    /// Produce a fixed-dimension embedding from text.
    fn embed(&self, text: &str) -> Vec<f32>;

    /// Output dimension of embeddings.
    fn dimension(&self) -> usize {
        EMBEDDING_DIM
    }

    /// Whether this model requires a runtime download or external binary.
    fn requires_external_runtime(&self) -> bool {
        false
    }
}

/// Hash-based embedder wrapping the existing `embedder::embed_text`.
///
/// This is the default when no neural model is available.
pub struct HashEmbedder;

impl EmbeddingModel for HashEmbedder {
    fn name(&self) -> &str {
        "hash-768d"
    }

    fn embed(&self, text: &str) -> Vec<f32> {
        embedder::embed_text(text)
    }

    fn dimension(&self) -> usize {
        embedder::EMBEDDING_DIM
    }
}

/// ONNX Runtime embedding model.
///
/// When the `onnx` feature flag is enabled, this struct loads an ONNX model
/// (e.g., CodeBERT) and a tokenizer from a directory, producing real neural
/// embeddings. Falls back to the hash embedder if model loading fails or
/// inference errors occur.
///
/// Expected model directory layout:
/// ```text
/// ~/.furnace/models/codebert/
///   model.onnx        # exported via optimum-cli or similar
///   tokenizer.json    # HuggingFace tokenizer config
/// ```
#[cfg(feature = "onnx")]
pub struct OnnxEmbedder {
    session: std::sync::Mutex<Option<ort::session::Session>>,
    tokenizer: Option<tokenizers::Tokenizer>,
    fallback: HashEmbedder,
    model_name: String,
    cached_dim: usize,
}

#[cfg(feature = "onnx")]
impl OnnxEmbedder {
    /// Maximum sequence length for tokenization (CodeBERT default).
    const MAX_SEQ_LEN: usize = 512;

    /// Create an ONNX embedder from a model directory.
    ///
    /// The directory must contain `model.onnx` and `tokenizer.json`.
    /// Falls back to hash embedder if either file is missing or loading fails.
    pub fn new(model_dir: &str) -> Self {
        let model_path = std::path::Path::new(model_dir).join("model.onnx");
        let tokenizer_path = std::path::Path::new(model_dir).join("tokenizer.json");

        let session = match ort::session::Session::builder()
            .and_then(|mut b| b.commit_from_file(&model_path))
        {
            Ok(s) => {
                tracing::info!(
                    model = %model_path.display(),
                    "ONNX model loaded successfully"
                );
                Some(s)
            }
            Err(e) => {
                tracing::warn!(
                    model = %model_path.display(),
                    error = %e,
                    "Failed to load ONNX model, falling back to hash embedder"
                );
                None
            }
        };

        let tokenizer = match tokenizers::Tokenizer::from_file(&tokenizer_path) {
            Ok(mut t) => {
                // Enable truncation to MAX_SEQ_LEN.
                let _ = t.with_truncation(Some(tokenizers::TruncationParams {
                    max_length: Self::MAX_SEQ_LEN,
                    ..Default::default()
                }));
                tracing::info!(
                    tokenizer = %tokenizer_path.display(),
                    "Tokenizer loaded successfully"
                );
                Some(t)
            }
            Err(e) => {
                tracing::warn!(
                    tokenizer = %tokenizer_path.display(),
                    error = %e,
                    "Failed to load tokenizer, falling back to hash embedder"
                );
                None
            }
        };

        let loaded = session.is_some() && tokenizer.is_some();

        let model_name = if loaded {
            format!("onnx-{}", std::path::Path::new(model_dir)
                .file_name()
                .and_then(|n| n.to_str())
                .unwrap_or("unknown"))
        } else {
            "hash-768d (onnx fallback)".to_string()
        };

        // Read hidden dim from output shape metadata before moving session into Mutex.
        let cached_dim = if let Some(ref s) = session {
            s.outputs()
                .first()
                .and_then(|o| {
                    if let ort::value::ValueType::Tensor { ref shape, .. } = *o.dtype() {
                        let dims: &[i64] = shape;
                        dims.last().copied().filter(|&d| d > 0).map(|d| d as usize)
                    } else {
                        None
                    }
                })
                .unwrap_or(EMBEDDING_DIM)
        } else {
            EMBEDDING_DIM
        };

        Self {
            session: std::sync::Mutex::new(session),
            tokenizer,
            fallback: HashEmbedder,
            model_name,
            cached_dim,
        }
    }

    /// Returns true if both model and tokenizer loaded successfully.
    pub fn is_loaded(&self) -> bool {
        let guard = self.session.lock().unwrap_or_else(|e| e.into_inner());
        guard.is_some() && self.tokenizer.is_some()
    }

    /// Run ONNX inference and return the embedding, or None on failure.
    fn infer(&self, text: &str) -> Option<Vec<f32>> {
        let tokenizer = self.tokenizer.as_ref()?;

        // 1. Tokenize
        let encoding = tokenizer.encode(text, true).ok()?;
        let ids = encoding.get_ids();
        let attention_mask = encoding.get_attention_mask();
        let seq_len = ids.len();

        if seq_len == 0 {
            return None;
        }

        // 2. Build input tensors via ort::value::Tensor
        let input_ids: Vec<i64> = ids.iter().map(|&id| id as i64).collect();
        let attn_mask: Vec<i64> = attention_mask.iter().map(|&m| m as i64).collect();

        let input_ids_arr = ndarray::Array2::from_shape_vec((1, seq_len), input_ids).ok()?;
        let attn_mask_arr = ndarray::Array2::from_shape_vec((1, seq_len), attn_mask).ok()?;

        let input_ids_tensor = ort::value::Tensor::from_array(input_ids_arr).ok()?;
        let attn_mask_tensor = ort::value::Tensor::from_array(attn_mask_arr).ok()?;

        // 3. Run the model (needs &mut Session)
        let mut session_guard = self.session.lock().unwrap_or_else(|e| e.into_inner());
        let session = session_guard.as_mut()?;

        let outputs = session
            .run(ort::inputs![input_ids_tensor, attn_mask_tensor])
            .ok()?;

        // 4. Extract last_hidden_state (first output, shape [1, seq_len, hidden_dim])
        let (shape, data) = outputs[0]
            .try_extract_tensor::<f32>()
            .ok()?;

        let shape_slice: &[i64] = shape;
        if shape_slice.len() != 3 || shape_slice[0] != 1 {
            tracing::warn!(shape = ?shape_slice, "Unexpected ONNX output shape");
            return None;
        }
        let hidden_dim = shape_slice[2] as usize;

        // 5. Mean-pool over sequence dimension, respecting attention mask
        let mut pooled = vec![0.0f32; hidden_dim];
        let mut mask_sum = 0.0f32;

        for t in 0..seq_len {
            let mask_val = attention_mask[t] as f32;
            mask_sum += mask_val;
            for d in 0..hidden_dim {
                pooled[d] += data[t * hidden_dim + d] * mask_val;
            }
        }

        if mask_sum > 0.0 {
            for d in 0..hidden_dim {
                pooled[d] /= mask_sum;
            }
        }

        // 6. L2 normalize
        let norm: f32 = pooled.iter().map(|x| x * x).sum::<f32>().sqrt();
        if norm > 1e-12 {
            for x in &mut pooled {
                *x /= norm;
            }
        }

        Some(pooled)
    }
}

// Safety: Session is protected by Mutex, Tokenizer is read-only after init.
#[cfg(feature = "onnx")]
unsafe impl Sync for OnnxEmbedder {}

#[cfg(feature = "onnx")]
impl EmbeddingModel for OnnxEmbedder {
    fn name(&self) -> &str {
        &self.model_name
    }

    fn embed(&self, text: &str) -> Vec<f32> {
        if self.is_loaded() {
            match self.infer(text) {
                Some(embedding) => embedding,
                None => {
                    tracing::debug!("ONNX inference failed, falling back to hash embedder");
                    self.fallback.embed(text)
                }
            }
        } else {
            self.fallback.embed(text)
        }
    }

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

    fn requires_external_runtime(&self) -> bool {
        true
    }
}

/// Create the best available embedding model.
///
/// Returns an ONNX model if the feature is enabled and a model path is
/// provided, otherwise returns the hash embedder.
pub fn create_embedder(_model_path: Option<&str>) -> Box<dyn EmbeddingModel> {
    #[cfg(feature = "onnx")]
    if let Some(path) = _model_path {
        return Box::new(OnnxEmbedder::new(path));
    }

    Box::new(HashEmbedder)
}

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

    #[test]
    fn test_hash_embedder_name() {
        let embedder = HashEmbedder;
        assert_eq!(embedder.name(), "hash-768d");
    }

    #[test]
    fn test_hash_embedder_dimension() {
        let embedder = HashEmbedder;
        assert_eq!(embedder.dimension(), EMBEDDING_DIM);
    }

    #[test]
    fn test_hash_embedder_produces_vectors() {
        let embedder = HashEmbedder;
        let vec = embedder.embed("hello world");
        assert_eq!(vec.len(), EMBEDDING_DIM);
        // Should be normalized.
        let norm: f32 = vec.iter().map(|x| x * x).sum::<f32>().sqrt();
        assert!((norm - 1.0).abs() < 1e-5);
    }

    #[test]
    fn test_hash_embedder_deterministic() {
        let embedder = HashEmbedder;
        let v1 = embedder.embed("test input");
        let v2 = embedder.embed("test input");
        assert_eq!(v1, v2);
    }

    #[test]
    fn test_hash_embedder_no_external_runtime() {
        let embedder = HashEmbedder;
        assert!(!embedder.requires_external_runtime());
    }

    #[test]
    fn test_create_embedder_default() {
        let embedder = create_embedder(None);
        assert_eq!(embedder.name(), "hash-768d");
        assert_eq!(embedder.dimension(), EMBEDDING_DIM);
    }

    #[test]
    fn test_embedding_trait_object() {
        // Ensure EmbeddingModel can be used as a trait object.
        let embedder: Box<dyn EmbeddingModel> = Box::new(HashEmbedder);
        let vec = embedder.embed("trait object test");
        assert_eq!(vec.len(), EMBEDDING_DIM);
    }
}
