cml-rs 0.4.0

//! SQLite-based embedding store with FTS5 hybrid search
//!
//! This module provides a high-performance embedding lookup table that combines:
//! - Full-text search (FTS5) for keyword matching
//! - Vector similarity search for semantic matching
//! - Hierarchical parent-child relationships
//!
//! The hybrid approach dramatically outperforms either method alone:
//! - FTS5 provides high precision on exact keywords
//! - Vector search provides high recall on semantic similarity
//! - Combined: Best of both worlds

use crate::chunker::Chunk;
use rusqlite::{params, Connection, Result as SqlResult};
use std::path::Path;

/// Dimension of embedding vectors (matches sentence-transformers MiniLM-L6-v2).
pub const EMBEDDING_DIM: usize = 384;

/// SQLite embedding store with FTS5 hybrid search
pub struct EmbeddingStore {
    conn: Connection,
}

impl EmbeddingStore {
    /// Create a new embedding store (in-memory for testing)
    pub fn new_in_memory() -> SqlResult<Self> {
        let conn = Connection::open_in_memory()?;
        Self::init_schema(&conn)?;
        Ok(Self { conn })
    }

    /// Open or create an embedding store from file
    pub fn open(path: &Path) -> SqlResult<Self> {
        let conn = Connection::open(path)?;
        Self::init_schema(&conn)?;
        Ok(Self { conn })
    }

    /// Initialize database schema
    fn init_schema(conn: &Connection) -> SqlResult<()> {
        // Main chunks table
        conn.execute(
            "CREATE TABLE IF NOT EXISTS chunks (
                id TEXT PRIMARY KEY,
                parent_id TEXT,
                content_hash TEXT NOT NULL,
                profile TEXT NOT NULL,
                element_type TEXT NOT NULL,
                content TEXT NOT NULL,
                token_count INTEGER NOT NULL,
                metadata JSON,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )",
            [],
        )?;

        // Embeddings table (384-dim f32 vectors as BLOB)
        conn.execute(
            "CREATE TABLE IF NOT EXISTS embeddings (
                chunk_id TEXT PRIMARY KEY,
                embedding BLOB NOT NULL,
                norm REAL NOT NULL,
                FOREIGN KEY (chunk_id) REFERENCES chunks(id) ON DELETE CASCADE
            )",
            [],
        )?;

        // FTS5 virtual table for full-text search
        conn.execute(
            "CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(
                id,
                content,
                element_type,
                metadata,
                content='chunks',
                content_rowid='rowid'
            )",
            [],
        )?;

        // Indexes for fast lookups
        conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_parent ON chunks(parent_id)",
            [],
        )?;

        conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_profile ON chunks(profile, element_type)",
            [],
        )?;

        conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_content_hash ON chunks(content_hash)",
            [],
        )?;

        // Trigger to keep FTS5 in sync
        conn.execute(
            "CREATE TRIGGER IF NOT EXISTS chunks_fts_insert AFTER INSERT ON chunks BEGIN
                INSERT INTO chunks_fts(rowid, id, content, element_type, metadata)
                VALUES (new.rowid, new.id, new.content, new.element_type, new.metadata);
            END",
            [],
        )?;

        conn.execute(
            "CREATE TRIGGER IF NOT EXISTS chunks_fts_delete AFTER DELETE ON chunks BEGIN
                DELETE FROM chunks_fts WHERE rowid = old.rowid;
            END",
            [],
        )?;

        conn.execute(
            "CREATE TRIGGER IF NOT EXISTS chunks_fts_update AFTER UPDATE ON chunks BEGIN
                UPDATE chunks_fts SET
                    id = new.id,
                    content = new.content,
                    element_type = new.element_type,
                    metadata = new.metadata
                WHERE rowid = new.rowid;
            END",
            [],
        )?;

        Ok(())
    }

    /// Insert a chunk with its embedding
    pub fn insert_chunk(&mut self, chunk: &Chunk, embedding: &[f32]) -> SqlResult<()> {
        if embedding.len() != EMBEDDING_DIM {
            return Err(rusqlite::Error::InvalidParameterCount(
                EMBEDDING_DIM,
                embedding.len(),
            ));
        }

        let metadata_json = serde_json::to_string(&chunk.metadata)
            .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?;

        // Insert chunk
        self.conn.execute(
            "INSERT INTO chunks (id, parent_id, content_hash, profile, element_type, content, token_count, metadata)
             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
            params![
                chunk.id,
                chunk.parent_id,
                chunk.content_hash,
                chunk.profile,
                chunk.element_type,
                chunk.content,
                chunk.token_count,
                metadata_json,
            ],
        )?;

        // Convert embedding to BLOB (4 bytes per f32)
        let embedding_blob = embedding
            .iter()
            .flat_map(|f| f.to_le_bytes())
            .collect::<Vec<u8>>();

        // Calculate L2 norm
        let norm = Self::l2_norm(embedding);

        // Insert embedding
        self.conn.execute(
            "INSERT INTO embeddings (chunk_id, embedding, norm) VALUES (?1, ?2, ?3)",
            params![chunk.id, embedding_blob, norm],
        )?;

        Ok(())
    }

    /// Get a chunk by ID
    pub fn get_chunk(&self, id: &str) -> SqlResult<Option<Chunk>> {
        let mut stmt = self.conn.prepare(
            "SELECT id, parent_id, content_hash, profile, element_type, content, token_count, metadata
             FROM chunks WHERE id = ?1",
        )?;

        let mut rows = stmt.query(params![id])?;

        if let Some(row) = rows.next()? {
            let metadata_json: String = row.get(7)?;
            let metadata = serde_json::from_str(&metadata_json).map_err(|e| {
                rusqlite::Error::FromSqlConversionFailure(
                    7,
                    rusqlite::types::Type::Text,
                    Box::new(e),
                )
            })?;

            Ok(Some(Chunk {
                id: row.get(0)?,
                parent_id: row.get(1)?,
                content_hash: row.get(2)?,
                profile: row.get(3)?,
                element_type: row.get(4)?,
                content: row.get(5)?,
                token_count: row.get(6)?,
                metadata,
            }))
        } else {
            Ok(None)
        }
    }

    /// Get embedding for a chunk
    pub fn get_embedding(&self, chunk_id: &str) -> SqlResult<Option<Vec<f32>>> {
        let mut stmt = self
            .conn
            .prepare("SELECT embedding FROM embeddings WHERE chunk_id = ?1")?;

        let mut rows = stmt.query(params![chunk_id])?;

        if let Some(row) = rows.next()? {
            let blob: Vec<u8> = row.get(0)?;
            let embedding = Self::blob_to_embedding(&blob)?;
            Ok(Some(embedding))
        } else {
            Ok(None)
        }
    }

    /// FTS5 keyword search
    pub fn search_keywords(&self, query: &str, limit: usize) -> SqlResult<Vec<ChunkMatch>> {
        let mut stmt = self.conn.prepare(
            "SELECT c.id, c.content, c.element_type, c.profile, rank
             FROM chunks_fts
             JOIN chunks c ON chunks_fts.rowid = c.rowid
             WHERE chunks_fts MATCH ?1
             ORDER BY rank
             LIMIT ?2",
        )?;

        let mut rows = stmt.query(params![query, limit as i64])?;
        let mut matches = Vec::new();

        while let Some(row) = rows.next()? {
            matches.push(ChunkMatch {
                id: row.get(0)?,
                content: row.get(1)?,
                element_type: row.get(2)?,
                profile: row.get(3)?,
                score: row.get::<_, f64>(4)? as f32,
                match_type: MatchType::Keyword,
            });
        }

        Ok(matches)
    }

    /// Vector similarity search (brute force for now, fast enough for <100K chunks)
    pub fn search_similar(
        &self,
        query_embedding: &[f32],
        limit: usize,
    ) -> SqlResult<Vec<ChunkMatch>> {
        if query_embedding.len() != EMBEDDING_DIM {
            return Err(rusqlite::Error::InvalidParameterCount(
                EMBEDDING_DIM,
                query_embedding.len(),
            ));
        }

        let query_norm = Self::l2_norm(query_embedding);

        let mut stmt = self.conn.prepare(
            "SELECT c.id, c.content, c.element_type, c.profile, e.embedding, e.norm
             FROM chunks c
             JOIN embeddings e ON c.id = e.chunk_id",
        )?;

        let mut rows = stmt.query([])?;
        let mut matches = Vec::new();

        while let Some(row) = rows.next()? {
            let id: String = row.get(0)?;
            let content: String = row.get(1)?;
            let element_type: String = row.get(2)?;
            let profile: String = row.get(3)?;
            let embedding_blob: Vec<u8> = row.get(4)?;
            let norm: f32 = row.get(5)?;

            let embedding = Self::blob_to_embedding(&embedding_blob)?;

            // Cosine similarity = dot product / (norm1 * norm2)
            let dot_product: f32 = query_embedding
                .iter()
                .zip(&embedding)
                .map(|(a, b)| a * b)
                .sum();

            let similarity = dot_product / (query_norm * norm);

            matches.push(ChunkMatch {
                id,
                content,
                element_type,
                profile,
                score: similarity,
                match_type: MatchType::Vector,
            });
        }

        // Sort by similarity (descending) and take top N
        matches.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap());
        matches.truncate(limit);

        Ok(matches)
    }

    /// Hybrid search: Combine FTS5 + vector similarity
    pub fn hybrid_search(
        &self,
        keywords: &str,
        query_embedding: &[f32],
        limit: usize,
    ) -> SqlResult<Vec<ChunkMatch>> {
        // Get keyword matches (precision)
        let keyword_matches = self.search_keywords(keywords, limit * 2)?;

        // Get vector matches (recall)
        let vector_matches = self.search_similar(query_embedding, limit * 2)?;

        // Combine and rerank
        let mut combined = Self::merge_and_rerank(keyword_matches, vector_matches);
        combined.truncate(limit);

        Ok(combined)
    }

    /// Merge keyword and vector matches, rerank by combined score
    fn merge_and_rerank(
        keyword_matches: Vec<ChunkMatch>,
        vector_matches: Vec<ChunkMatch>,
    ) -> Vec<ChunkMatch> {
        use std::collections::HashMap;

        let mut matches_by_id: HashMap<String, ChunkMatch> = HashMap::new();
        let mut scores: HashMap<String, (f32, f32)> = HashMap::new(); // (keyword_score, vector_score)

        // Collect keyword scores and matches
        for m in keyword_matches {
            scores.entry(m.id.clone()).or_insert((0.0, 0.0)).0 = m.score.abs(); // FTS5 rank is negative
            matches_by_id.insert(m.id.clone(), m);
        }

        // Collect vector scores and matches
        for m in vector_matches {
            scores.entry(m.id.clone()).or_insert((0.0, 0.0)).1 = m.score;
            matches_by_id.entry(m.id.clone()).or_insert(m);
        }

        // Rerank: combined_score = 0.3 * keyword + 0.7 * vector (favor semantic)
        let mut combined: Vec<_> = scores
            .into_iter()
            .filter_map(|(id, (kw_score, vec_score))| {
                let combined_score = 0.3 * kw_score + 0.7 * vec_score;
                matches_by_id.get(&id).map(|m| {
                    let mut new_match = m.clone();
                    new_match.score = combined_score;
                    new_match.match_type = MatchType::Hybrid;
                    new_match
                })
            })
            .collect();

        combined.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap());
        combined
    }

    /// Get all child chunks of a parent
    pub fn get_children(&self, parent_id: &str) -> SqlResult<Vec<Chunk>> {
        let mut stmt = self.conn.prepare(
            "SELECT id, parent_id, content_hash, profile, element_type, content, token_count, metadata
             FROM chunks WHERE parent_id = ?1
             ORDER BY id",
        )?;

        let mut rows = stmt.query(params![parent_id])?;
        let mut children = Vec::new();

        while let Some(row) = rows.next()? {
            let metadata_json: String = row.get(7)?;
            let metadata = serde_json::from_str(&metadata_json).map_err(|e| {
                rusqlite::Error::FromSqlConversionFailure(
                    7,
                    rusqlite::types::Type::Text,
                    Box::new(e),
                )
            })?;

            children.push(Chunk {
                id: row.get(0)?,
                parent_id: row.get(1)?,
                content_hash: row.get(2)?,
                profile: row.get(3)?,
                element_type: row.get(4)?,
                content: row.get(5)?,
                token_count: row.get(6)?,
                metadata,
            });
        }

        Ok(children)
    }

    /// Count total chunks
    pub fn count_chunks(&self) -> SqlResult<usize> {
        let count: i64 = self
            .conn
            .query_row("SELECT COUNT(*) FROM chunks", [], |row| row.get(0))?;
        Ok(count as usize)
    }

    /// Calculate L2 norm of a vector
    fn l2_norm(vec: &[f32]) -> f32 {
        vec.iter().map(|x| x * x).sum::<f32>().sqrt()
    }

    /// Convert BLOB to f32 vector
    fn blob_to_embedding(blob: &[u8]) -> SqlResult<Vec<f32>> {
        if blob.len() != EMBEDDING_DIM * 4 {
            return Err(rusqlite::Error::InvalidColumnType(
                0,
                "Embedding BLOB".to_string(),
                rusqlite::types::Type::Blob,
            ));
        }

        let embedding = blob
            .chunks_exact(4)
            .map(|chunk| {
                let bytes = [chunk[0], chunk[1], chunk[2], chunk[3]];
                f32::from_le_bytes(bytes)
            })
            .collect();

        Ok(embedding)
    }
}

/// Search result match
#[derive(Debug, Clone, PartialEq)]
pub struct ChunkMatch {
    pub id: String,
    pub content: String,
    pub element_type: String,
    pub profile: String,
    pub score: f32,
    pub match_type: MatchType,
}

/// Type of search match
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MatchType {
    Keyword,
    Vector,
    Hybrid,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::id_generator::ElementId;
    use std::collections::HashMap;

    fn create_test_chunk(id: &str, content: &str) -> Chunk {
        Chunk {
            id: id.to_string(),
            parent_id: None,
            content_hash: ElementId::new(id, content).content_hash,
            profile: "code:api".to_string(),
            element_type: "function".to_string(),
            content: content.to_string(),
            token_count: content.len() / 4,
            metadata: HashMap::new(),
        }
    }

    fn create_test_embedding() -> Vec<f32> {
        vec![0.1; EMBEDDING_DIM]
    }

    #[test]
    fn test_create_store() {
        let store = EmbeddingStore::new_in_memory();
        assert!(store.is_ok());
    }

    #[test]
    fn test_insert_and_get_chunk() {
        let mut store = EmbeddingStore::new_in_memory().unwrap();
        let chunk = create_test_chunk("test.id", "Test content");
        let embedding = create_test_embedding();

        store.insert_chunk(&chunk, &embedding).unwrap();

        let retrieved = store.get_chunk("test.id").unwrap();
        assert!(retrieved.is_some());
        assert_eq!(retrieved.unwrap().content, "Test content");
    }

    #[test]
    fn test_get_embedding() {
        let mut store = EmbeddingStore::new_in_memory().unwrap();
        let chunk = create_test_chunk("test.id", "Test content");
        let embedding = create_test_embedding();

        store.insert_chunk(&chunk, &embedding).unwrap();

        let retrieved_emb = store.get_embedding("test.id").unwrap();
        assert!(retrieved_emb.is_some());
        assert_eq!(retrieved_emb.unwrap().len(), EMBEDDING_DIM);
    }

    #[test]
    fn test_fts_search() {
        let mut store = EmbeddingStore::new_in_memory().unwrap();

        let chunk1 = create_test_chunk("test.1", "Vector push method");
        let chunk2 = create_test_chunk("test.2", "HashMap insert function");
        let embedding = create_test_embedding();

        store.insert_chunk(&chunk1, &embedding).unwrap();
        store.insert_chunk(&chunk2, &embedding).unwrap();

        let results = store.search_keywords("vector", 10).unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].id, "test.1");
    }

    #[test]
    fn test_vector_similarity() {
        let mut store = EmbeddingStore::new_in_memory().unwrap();

        let chunk = create_test_chunk("test.id", "Test content");
        let embedding = create_test_embedding();

        store.insert_chunk(&chunk, &embedding).unwrap();

        // Query with same embedding should have similarity ~1.0
        let results = store.search_similar(&embedding, 10).unwrap();
        assert_eq!(results.len(), 1);
        assert!((results[0].score - 1.0).abs() < 0.01);
    }

    #[test]
    fn test_hybrid_search() {
        let mut store = EmbeddingStore::new_in_memory().unwrap();

        let chunk1 = create_test_chunk("test.1", "Vector push method adds items");
        let chunk2 = create_test_chunk("test.2", "HashMap insert stores key-value pairs");
        let embedding = create_test_embedding();

        store.insert_chunk(&chunk1, &embedding).unwrap();
        store.insert_chunk(&chunk2, &embedding).unwrap();

        let results = store.hybrid_search("vector", &embedding, 10).unwrap();
        assert!(results.len() > 0);
        assert_eq!(results[0].match_type, MatchType::Hybrid);
    }

    #[test]
    fn test_parent_child_relationship() {
        let mut store = EmbeddingStore::new_in_memory().unwrap();

        let parent = create_test_chunk("parent.id", "Parent content");
        let mut child = create_test_chunk("parent.id#0", "Child content");
        child.parent_id = Some("parent.id".to_string());

        let embedding = create_test_embedding();

        store.insert_chunk(&parent, &embedding).unwrap();
        store.insert_chunk(&child, &embedding).unwrap();

        let children = store.get_children("parent.id").unwrap();
        assert_eq!(children.len(), 1);
        assert_eq!(children[0].id, "parent.id#0");
    }

    #[test]
    fn test_count_chunks() {
        let mut store = EmbeddingStore::new_in_memory().unwrap();
        let embedding = create_test_embedding();

        assert_eq!(store.count_chunks().unwrap(), 0);

        store
            .insert_chunk(&create_test_chunk("test.1", "Content 1"), &embedding)
            .unwrap();
        store
            .insert_chunk(&create_test_chunk("test.2", "Content 2"), &embedding)
            .unwrap();

        assert_eq!(store.count_chunks().unwrap(), 2);
    }
}