jina_cache.rs

//! Jina Embedding Cache with Sparse API Usage
//!
//! Strategy:
//! 1. Exact match in HashMap → use cached (0 API calls)
//! 2. Near match (Hamming < 0.15) → use closest cached (0 API calls)  
//! 3. Cache miss → call Jina API, then cache result
//!
//! For typical knowledge graphs with repeated entities,
//! this reduces Jina API calls by 90%+

use std::collections::HashMap;
use std::fs::{File, OpenOptions};
use std::io::{BufReader, BufWriter, Read, Write};
use std::path::Path;

// Same fingerprint structure as main.rs
const N: usize = 10_000;
const N64: usize = 157;
const NEAR_THRESHOLD: u32 = 1500;  // 0.15 * 10000 = 15% Hamming distance

#[repr(align(64))]
#[derive(Clone)]
pub struct Fingerprint {
    pub data: [u64; N64],
}

impl Fingerprint {
    pub fn zero() -> Self { Self { data: [0u64; N64] } }
    
    #[inline]
    pub fn hamming(&self, other: &Fingerprint) -> u32 {
        let mut t = 0u32;
        for i in 0..N64 { t += (self.data[i] ^ other.data[i]).count_ones(); }
        t
    }
    
    pub fn similarity(&self, other: &Fingerprint) -> f64 {
        1.0 - (self.hamming(other) as f64 / N as f64)
    }
    
    /// Convert from f32 Jina embedding (1024D) to binary fingerprint (10Kbit)
    pub fn from_jina_embedding(embedding: &[f32]) -> Self {
        let mut fp = Fingerprint::zero();
        
        // Method: threshold at median, then expand to 10K bits
        // Each of 1024 dimensions maps to ~10 bits
        let mut sorted: Vec<f32> = embedding.to_vec();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
        let median = sorted[512];
        
        for (i, &val) in embedding.iter().enumerate() {
            let base_bit = i * 10;  // 1024 * 10 = 10240 > 10000, so we wrap
            
            // Set multiple bits based on value relative to median
            let strength = ((val - median).abs() * 5.0).min(5.0) as usize;
            
            for j in 0..strength {
                let bit_pos = (base_bit + j) % N;
                let word_idx = bit_pos / 64;
                let bit_idx = bit_pos % 64;
                
                if val > median {
                    fp.data[word_idx] |= 1 << bit_idx;
                }
            }
        }
        
        fp
    }
    
    /// Serialize to bytes
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::with_capacity(N64 * 8);
        for word in &self.data {
            bytes.extend_from_slice(&word.to_le_bytes());
        }
        bytes
    }
    
    /// Deserialize from bytes
    pub fn from_bytes(bytes: &[u8]) -> Option<Self> {
        if bytes.len() != N64 * 8 { return None; }
        let mut fp = Fingerprint::zero();
        for (i, chunk) in bytes.chunks(8).enumerate() {
            fp.data[i] = u64::from_le_bytes(chunk.try_into().ok()?);
        }
        Some(fp)
    }
}

/// Cache entry with original text and fingerprint
#[derive(Clone)]
struct CacheEntry {
    text: String,
    fingerprint: Fingerprint,
    jina_embedding: Option<Vec<f32>>,  // Keep original for precision if needed
}

/// Jina embedding cache with sparse API usage
pub struct JinaCache {
    /// Exact match lookup
    exact: HashMap<String, CacheEntry>,
    
    /// All entries for near-match search (could use a proper ANN index)
    entries: Vec<CacheEntry>,
    
    /// API key
    api_key: String,
    
    /// Statistics
    pub stats: CacheStats,
    
    /// Persistence path
    cache_path: Option<String>,
}

#[derive(Default, Clone)]
pub struct CacheStats {
    pub exact_hits: u64,
    pub near_hits: u64,
    pub api_calls: u64,
    pub total_lookups: u64,
}

impl CacheStats {
    pub fn hit_rate(&self) -> f64 {
        if self.total_lookups == 0 { return 0.0; }
        (self.exact_hits + self.near_hits) as f64 / self.total_lookups as f64
    }
    
    pub fn api_call_rate(&self) -> f64 {
        if self.total_lookups == 0 { return 0.0; }
        self.api_calls as f64 / self.total_lookups as f64
    }
}

impl JinaCache {
    pub fn new(api_key: &str) -> Self {
        Self {
            exact: HashMap::new(),
            entries: Vec::new(),
            api_key: api_key.to_string(),
            stats: CacheStats::default(),
            cache_path: None,
        }
    }
    
    pub fn with_persistence(mut self, path: &str) -> Self {
        self.cache_path = Some(path.to_string());
        self.load_from_disk();
        self
    }
    
    /// Get or create fingerprint for text
    pub fn get_fingerprint(&mut self, text: &str) -> Result<Fingerprint, String> {
        self.stats.total_lookups += 1;
        
        // 1. Exact match
        if let Some(entry) = self.exact.get(text) {
            self.stats.exact_hits += 1;
            return Ok(entry.fingerprint.clone());
        }
        
        // 2. Near match (linear scan - could use ANN for large caches)
        let query_lower = text.to_lowercase();
        for entry in &self.entries {
            // Quick string similarity check first
            if string_similar(&entry.text, text) {
                self.stats.near_hits += 1;
                return Ok(entry.fingerprint.clone());
            }
        }
        
        // 3. API call needed
        self.stats.api_calls += 1;
        let embedding = self.call_jina_api(text)?;
        let fingerprint = Fingerprint::from_jina_embedding(&embedding);
        
        // Cache it
        let entry = CacheEntry {
            text: text.to_string(),
            fingerprint: fingerprint.clone(),
            jina_embedding: Some(embedding),
        };
        
        self.exact.insert(text.to_string(), entry.clone());
        self.entries.push(entry);
        
        // Persist
        if self.cache_path.is_some() {
            self.save_to_disk();
        }
        
        Ok(fingerprint)
    }
    
    /// Batch get - more efficient for multiple texts
    pub fn get_fingerprints_batch(&mut self, texts: &[&str]) -> Result<Vec<Fingerprint>, String> {
        let mut results = Vec::with_capacity(texts.len());
        let mut to_fetch: Vec<(usize, &str)> = Vec::new();
        
        // Check cache first
        for (i, text) in texts.iter().enumerate() {
            self.stats.total_lookups += 1;
            
            if let Some(entry) = self.exact.get(*text) {
                self.stats.exact_hits += 1;
                results.push((i, entry.fingerprint.clone()));
            } else {
                // Check near matches
                let mut found = false;
                for entry in &self.entries {
                    if string_similar(&entry.text, text) {
                        self.stats.near_hits += 1;
                        results.push((i, entry.fingerprint.clone()));
                        found = true;
                        break;
                    }
                }
                if !found {
                    to_fetch.push((i, *text));
                }
            }
        }
        
        // Batch API call for misses
        if !to_fetch.is_empty() {
            let texts_to_fetch: Vec<&str> = to_fetch.iter().map(|(_, t)| *t).collect();
            let embeddings = self.call_jina_api_batch(&texts_to_fetch)?;
            
            for ((i, text), embedding) in to_fetch.into_iter().zip(embeddings.into_iter()) {
                self.stats.api_calls += 1;
                let fingerprint = Fingerprint::from_jina_embedding(&embedding);
                
                let entry = CacheEntry {
                    text: text.to_string(),
                    fingerprint: fingerprint.clone(),
                    jina_embedding: Some(embedding),
                };
                
                self.exact.insert(text.to_string(), entry.clone());
                self.entries.push(entry);
                results.push((i, fingerprint));
            }
        }
        
        // Sort by original index
        results.sort_by_key(|(i, _)| *i);
        Ok(results.into_iter().map(|(_, fp)| fp).collect())
    }
    
    /// Find near matches in cache (for debugging/analysis)
    pub fn find_near_matches(&self, text: &str, threshold: f64) -> Vec<(String, f64)> {
        let mut matches = Vec::new();
        
        // Get fingerprint for query (without caching)
        if let Some(entry) = self.exact.get(text) {
            for other in &self.entries {
                let sim = entry.fingerprint.similarity(&other.fingerprint);
                if sim >= threshold && other.text != text {
                    matches.push((other.text.clone(), sim));
                }
            }
        }
        
        matches.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
        matches
    }
    
    fn call_jina_api(&self, text: &str) -> Result<Vec<f32>, String> {
        // Placeholder - implement actual API call
        // For now, generate deterministic pseudo-embedding
        Ok(pseudo_embedding(text))
    }
    
    fn call_jina_api_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, String> {
        // Placeholder - implement actual batch API call
        Ok(texts.iter().map(|t| pseudo_embedding(t)).collect())
    }
    
    fn save_to_disk(&self) {
        if let Some(path) = self.cache_path.as_ref() {
            if let Ok(file) = File::create(path) {
                let mut writer = BufWriter::new(file);
                
                // Simple format: count, then (text_len, text, fingerprint_bytes) for each
                let count = self.entries.len() as u32;
                let _ = writer.write_all(&count.to_le_bytes());
                
                for entry in &self.entries {
                    let text_bytes = entry.text.as_bytes();
                    let text_len = text_bytes.len() as u32;
                    let _ = writer.write_all(&text_len.to_le_bytes());
                    let _ = writer.write_all(text_bytes);
                    let _ = writer.write_all(&entry.fingerprint.to_bytes());
                }
            }
        }
    }
    
    fn load_from_disk(&mut self) {
        if let Some(path) = self.cache_path.as_ref() {
            if let Ok(file) = File::open(path) {
                let mut reader = BufReader::new(file);
                
                let mut count_bytes = [0u8; 4];
                if reader.read_exact(&mut count_bytes).is_err() { return; }
                let count = u32::from_le_bytes(count_bytes) as usize;
                
                for _ in 0..count {
                    let mut len_bytes = [0u8; 4];
                    if reader.read_exact(&mut len_bytes).is_err() { break; }
                    let text_len = u32::from_le_bytes(len_bytes) as usize;
                    
                    let mut text_bytes = vec![0u8; text_len];
                    if reader.read_exact(&mut text_bytes).is_err() { break; }
                    let text = String::from_utf8_lossy(&text_bytes).to_string();
                    
                    let mut fp_bytes = vec![0u8; N64 * 8];
                    if reader.read_exact(&mut fp_bytes).is_err() { break; }
                    
                    if let Some(fingerprint) = Fingerprint::from_bytes(&fp_bytes) {
                        let entry = CacheEntry {
                            text: text.clone(),
                            fingerprint,
                            jina_embedding: None,
                        };
                        self.exact.insert(text, entry.clone());
                        self.entries.push(entry);
                    }
                }
            }
        }
    }
    
    pub fn len(&self) -> usize { self.entries.len() }
    
    pub fn print_stats(&self) {
        println!("JinaCache Statistics:");
        println!("  Entries:      {}", self.entries.len());
        println!("  Lookups:      {}", self.stats.total_lookups);
        println!("  Exact hits:   {} ({:.1}%)", 
                 self.stats.exact_hits,
                 100.0 * self.stats.exact_hits as f64 / self.stats.total_lookups.max(1) as f64);
        println!("  Near hits:    {} ({:.1}%)",
                 self.stats.near_hits,
                 100.0 * self.stats.near_hits as f64 / self.stats.total_lookups.max(1) as f64);
        println!("  API calls:    {} ({:.1}%)",
                 self.stats.api_calls,
                 100.0 * self.stats.api_calls as f64 / self.stats.total_lookups.max(1) as f64);
        println!("  Hit rate:     {:.1}%", 100.0 * self.stats.hit_rate());
    }
}

/// Simple string similarity for near-match detection
fn string_similar(a: &str, b: &str) -> bool {
    let a_lower = a.to_lowercase();
    let b_lower = b.to_lowercase();
    
    // Exact case-insensitive
    if a_lower == b_lower { return true; }
    
    // One is prefix/suffix of other
    if a_lower.starts_with(&b_lower) || b_lower.starts_with(&a_lower) { return true; }
    
    // Levenshtein distance <= 2 for short strings
    if a.len() <= 10 && b.len() <= 10 {
        if levenshtein(&a_lower, &b_lower) <= 2 { return true; }
    }
    
    false
}

fn levenshtein(a: &str, b: &str) -> usize {
    let a: Vec<char> = a.chars().collect();
    let b: Vec<char> = b.chars().collect();
    
    let mut dp = vec![vec![0usize; b.len() + 1]; a.len() + 1];
    
    for i in 0..=a.len() { dp[i][0] = i; }
    for j in 0..=b.len() { dp[0][j] = j; }
    
    for i in 1..=a.len() {
        for j in 1..=b.len() {
            let cost = if a[i-1] == b[j-1] { 0 } else { 1 };
            dp[i][j] = (dp[i-1][j] + 1)
                .min(dp[i][j-1] + 1)
                .min(dp[i-1][j-1] + cost);
        }
    }
    
    dp[a.len()][b.len()]
}

/// Pseudo-embedding for testing (replace with actual Jina API call)
fn pseudo_embedding(text: &str) -> Vec<f32> {
    use std::collections::hash_map::DefaultHasher;
    use std::hash::{Hash, Hasher};
    
    let mut embedding = vec![0.0f32; 1024];
    
    // Deterministic pseudo-random based on text
    for (i, chunk) in text.as_bytes().chunks(4).enumerate() {
        let mut hasher = DefaultHasher::new();
        chunk.hash(&mut hasher);
        i.hash(&mut hasher);
        let h = hasher.finish();
        
        for j in 0..8 {
            let idx = (i * 8 + j) % 1024;
            let val = ((h >> (j * 8)) & 0xFF) as f32 / 255.0 - 0.5;
            embedding[idx] += val;
        }
    }
    
    // Normalize
    let norm: f32 = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm > 0.0 {
        for x in &mut embedding { *x /= norm; }
    }
    
    embedding
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_cache_hit_rate() {
        let mut cache = JinaCache::new("test_key");

        // First access - all API calls (use distinct strings to avoid near-match)
        let texts = vec!["Ada", "Jan", "loves", "creates", "butterfly"];
        for text in &texts {
            let _ = cache.get_fingerprint(text);
        }

        assert_eq!(cache.stats.api_calls, 5);
        assert_eq!(cache.stats.exact_hits, 0);
        
        // Second access - all cache hits
        for text in &texts {
            let _ = cache.get_fingerprint(text);
        }
        
        assert_eq!(cache.stats.api_calls, 5);  // No new API calls
        assert_eq!(cache.stats.exact_hits, 5);
        
        println!("Hit rate: {:.1}%", 100.0 * cache.stats.hit_rate());
    }
    
    #[test]
    fn test_near_match() {
        let mut cache = JinaCache::new("test_key");
        
        // Cache "Ada"
        let _ = cache.get_fingerprint("Ada");
        
        // "ada" should near-match (case insensitive)
        let _ = cache.get_fingerprint("ada");
        
        assert_eq!(cache.stats.near_hits, 1);
        assert_eq!(cache.stats.api_calls, 1);  // Only one API call
    }
}