//! Text embedder for producing 768-dimensional vectors from text.
//!
//! Uses a word-level hashing approach with positional encoding and
//! subword decomposition. This is not a neural embedding model, but
//! produces meaningfully better vectors than character-frequency hashing
//! by capturing word-level and subword-level features.
//!
//! Both the search tier (query embedding) and the indexer (document
//! embedding) share this implementation for consistency.

use std::collections::HashMap;

/// Embedding dimension — matches common transformer model output sizes.
pub const EMBEDDING_DIM: usize = 768;

/// Produce a 768-dimensional embedding from text.
///
/// The approach:
/// 1. Tokenize into words (lowercased, stripped of punctuation)
/// 2. Hash each word and its character-level n-grams into buckets
/// 3. Apply IDF-like weighting (shorter, rarer-looking words get more weight)
/// 4. Add positional encoding for word order sensitivity
/// 5. L2-normalize the result
pub fn embed_text(text: &str) -> Vec<f32> {
    let mut vec = vec![0.0f32; EMBEDDING_DIM];
    let words = tokenize(text);

    if words.is_empty() {
        return vec;
    }

    // Track word frequencies for TF weighting.
    let mut word_freq: HashMap<&str, usize> = HashMap::new();
    for w in &words {
        *word_freq.entry(w.as_str()).or_insert(0) += 1;
    }
    let total_words = words.len() as f32;

    for (pos, word) in words.iter().enumerate() {
        let tf = (*word_freq.get(word.as_str()).unwrap_or(&1) as f32) / total_words;

        // IDF-like weight: shorter common words (the, a, is) get less weight.
        // Longer, more specific words get more weight.
        let idf = (1.0 + word.len() as f32).ln();

        let weight = tf * idf;

        // 1. Hash the full word into multiple dimensions.
        let h1 = hash_str(word, 0);
        let h2 = hash_str(word, 1);
        let h3 = hash_str(word, 2);
        vec[h1 % EMBEDDING_DIM] += weight;
        vec[h2 % EMBEDDING_DIM] += weight * 0.7;
        vec[h3 % EMBEDDING_DIM] += weight * 0.5;

        // 2. Character trigram hashing (captures subword structure).
        //    E.g., "function" → "fun", "unc", "nct", "cti", "tio", "ion"
        let word_bytes = word.as_bytes();
        for window in word_bytes.windows(3) {
            let h = trigram_hash(window);
            vec[h % EMBEDDING_DIM] += weight * 0.3;
        }

        // 3. Positional encoding: early words (likely more important) get
        //    a slight boost in different dimensions.
        let pos_weight = 1.0 / (1.0 + pos as f32 * 0.1);
        let pos_hash = hash_str(word, 3).wrapping_add(pos * 7);
        vec[pos_hash % EMBEDDING_DIM] += weight * pos_weight * 0.2;
    }

    // 4. Add bigram features (word pairs) for phrase sensitivity.
    for pair in words.windows(2) {
        let bigram = format!("{}_{}", pair[0], pair[1]);
        let h = hash_str(&bigram, 4);
        vec[h % EMBEDDING_DIM] += 0.5;
    }

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

    vec
}

/// Tokenize text into lowercase words, stripping punctuation.
fn tokenize(text: &str) -> Vec<String> {
    text.split(|c: char| c.is_whitespace() || c == '_' || c == '-' || c == '/' || c == '\\')
        .map(|w| {
            w.chars()
                .filter(|c| c.is_alphanumeric())
                .collect::<String>()
                .to_lowercase()
        })
        .filter(|w| !w.is_empty())
        .collect()
}

/// FNV-1a inspired hash with a seed for producing multiple independent hashes.
fn hash_str(s: &str, seed: usize) -> usize {
    let mut hash: usize = 0xcbf29ce484222325_usize.wrapping_add(seed.wrapping_mul(0x100000001b3));
    for b in s.bytes() {
        hash ^= b as usize;
        hash = hash.wrapping_mul(0x100000001b3);
    }
    hash
}

/// Hash a character trigram.
fn trigram_hash(window: &[u8]) -> usize {
    let mut h: usize = 0x811c9dc5;
    for &b in window {
        h ^= b as usize;
        h = h.wrapping_mul(0x01000193);
    }
    h
}

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

    #[test]
    fn test_embed_deterministic() {
        let e1 = embed_text("hello world");
        let e2 = embed_text("hello world");
        assert_eq!(e1, e2);
    }

    #[test]
    fn test_embed_normalized() {
        let e = embed_text("search for function definitions in Rust");
        let norm: f32 = e.iter().map(|x| x * x).sum::<f32>().sqrt();
        assert!((norm - 1.0).abs() < 1e-5, "norm was {}", norm);
    }

    #[test]
    fn test_embed_dimension() {
        let e = embed_text("anything");
        assert_eq!(e.len(), EMBEDDING_DIM);
    }

    #[test]
    fn test_embed_empty() {
        let e = embed_text("");
        assert_eq!(e.len(), EMBEDDING_DIM);
        assert!(e.iter().all(|&v| v == 0.0));
    }

    #[test]
    fn test_similar_texts_closer_than_different() {
        let e_rust_fn = embed_text("rust function definition");
        let e_rust_method = embed_text("rust method definition");
        let e_cooking = embed_text("cooking recipe for pasta");

        let sim_similar = cosine_similarity(&e_rust_fn, &e_rust_method);
        let sim_different = cosine_similarity(&e_rust_fn, &e_cooking);

        assert!(
            sim_similar > sim_different,
            "similar texts should have higher cosine similarity: {} vs {}",
            sim_similar,
            sim_different
        );
    }

    #[test]
    fn test_tokenize() {
        let tokens = tokenize("Hello_World-Test/Path");
        assert_eq!(tokens, vec!["hello", "world", "test", "path"]);
    }

    fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
        a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
    }
}
