rlm_rs/storage/

sqlite.rs

//! `SQLite` storage implementation.
//!
//! Provides persistent storage using `SQLite` with proper transaction
//! management and migration support.

// SQLite stores all integers as i64. These casts are intentional and safe
// because we only store non-negative values that fit in usize.
#![allow(clippy::cast_possible_truncation)]
#![allow(clippy::cast_sign_loss)]

use crate::core::{Buffer, BufferMetadata, Chunk, ChunkMetadata, Context};
use crate::error::{Result, StorageError};
use crate::storage::schema::{
    CHECK_SCHEMA_SQL, CURRENT_SCHEMA_VERSION, GET_VERSION_SQL, SCHEMA_SQL, SET_VERSION_SQL,
};
use crate::storage::traits::{Storage, StorageStats};
use rusqlite::{Connection, OptionalExtension, params};
use std::path::{Path, PathBuf};

/// SQLite-based storage implementation.
///
/// Provides persistent storage for RLM state with full ACID guarantees.
///
/// # Examples
///
/// ```no_run
/// use rlm_rs::storage::{SqliteStorage, Storage};
///
/// let mut storage = SqliteStorage::open("rlm-state.db").unwrap();
/// storage.init().unwrap();
/// ```
pub struct SqliteStorage {
    /// `SQLite` connection.
    conn: Connection,
    /// Path to the database file (None for in-memory).
    path: Option<PathBuf>,
}

impl SqliteStorage {
    /// Opens or creates a `SQLite` database at the given path.
    ///
    /// # Arguments
    ///
    /// * `path` - Path to the database file. Parent directory must exist.
    ///
    /// # Errors
    ///
    /// Returns an error if the database cannot be opened or initialized.
    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
        let path = path.as_ref().to_path_buf();

        // Ensure parent directory exists
        if let Some(parent) = path.parent()
            && !parent.exists()
        {
            std::fs::create_dir_all(parent).map_err(|e| StorageError::Database(e.to_string()))?;
        }

        let conn = Connection::open(&path).map_err(StorageError::from)?;

        // Enable foreign keys
        conn.execute("PRAGMA foreign_keys = ON;", [])
            .map_err(StorageError::from)?;

        // Use WAL mode for better concurrent access (returns result, use query_row)
        let _: String = conn
            .query_row("PRAGMA journal_mode = WAL;", [], |row| row.get(0))
            .map_err(StorageError::from)?;

        Ok(Self {
            conn,
            path: Some(path),
        })
    }

    /// Creates an in-memory `SQLite` database.
    ///
    /// Useful for testing.
    ///
    /// # Errors
    ///
    /// Returns an error if the database cannot be created.
    pub fn in_memory() -> Result<Self> {
        let conn = Connection::open_in_memory().map_err(StorageError::from)?;
        conn.execute("PRAGMA foreign_keys = ON;", [])
            .map_err(StorageError::from)?;

        Ok(Self { conn, path: None })
    }

    /// Returns the database path (None for in-memory).
    #[must_use]
    pub fn path(&self) -> Option<&Path> {
        self.path.as_deref()
    }

    /// Gets the current schema version.
    fn get_schema_version(&self) -> Result<Option<u32>> {
        let version: Option<String> = self
            .conn
            .query_row(GET_VERSION_SQL, [], |row| row.get(0))
            .optional()
            .map_err(StorageError::from)?;

        Ok(version.and_then(|v| v.parse().ok()))
    }

    /// Sets the schema version.
    fn set_schema_version(&self, version: u32) -> Result<()> {
        self.conn
            .execute(SET_VERSION_SQL, params![version.to_string()])
            .map_err(StorageError::from)?;
        Ok(())
    }

    /// Returns current Unix timestamp.
    #[allow(clippy::cast_possible_wrap)]
    fn now() -> i64 {
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map_or(0, |d| d.as_secs() as i64)
    }
}

impl Storage for SqliteStorage {
    fn init(&mut self) -> Result<()> {
        // Check if already initialized
        let is_init: i64 = self
            .conn
            .query_row(CHECK_SCHEMA_SQL, [], |row| row.get(0))
            .map_err(StorageError::from)?;

        if is_init == 0 {
            // Fresh install - create schema
            self.conn
                .execute_batch(SCHEMA_SQL)
                .map_err(StorageError::from)?;
            self.set_schema_version(CURRENT_SCHEMA_VERSION)?;
        } else if let Some(current) = self.get_schema_version()?
            && current < CURRENT_SCHEMA_VERSION
        {
            // Run migrations
            let migrations = crate::storage::schema::get_migrations_from(current);
            for migration in migrations {
                self.conn
                    .execute_batch(migration.sql)
                    .map_err(|e| StorageError::Migration(e.to_string()))?;
            }
            self.set_schema_version(CURRENT_SCHEMA_VERSION)?;
        }

        Ok(())
    }

    fn is_initialized(&self) -> Result<bool> {
        let count: i64 = self
            .conn
            .query_row(CHECK_SCHEMA_SQL, [], |row| row.get(0))
            .map_err(StorageError::from)?;
        Ok(count > 0)
    }

    fn reset(&mut self) -> Result<()> {
        self.conn
            .execute_batch(
                r"
            DELETE FROM chunk_embeddings;
            DELETE FROM chunks;
            DELETE FROM buffers;
            DELETE FROM context;
            DELETE FROM metadata;
        ",
            )
            .map_err(StorageError::from)?;
        Ok(())
    }

    // ==================== Context Operations ====================

    fn save_context(&mut self, context: &Context) -> Result<()> {
        let data = serde_json::to_string(context).map_err(StorageError::from)?;
        let now = Self::now();

        self.conn
            .execute(
                r"
            INSERT OR REPLACE INTO context (id, data, created_at, updated_at)
            VALUES (1, ?, COALESCE((SELECT created_at FROM context WHERE id = 1), ?), ?)
        ",
                params![data, now, now],
            )
            .map_err(StorageError::from)?;

        Ok(())
    }

    fn load_context(&self) -> Result<Option<Context>> {
        let data: Option<String> = self
            .conn
            .query_row("SELECT data FROM context WHERE id = 1", [], |row| {
                row.get(0)
            })
            .optional()
            .map_err(StorageError::from)?;

        match data {
            Some(json) => {
                let context = serde_json::from_str(&json).map_err(StorageError::from)?;
                Ok(Some(context))
            }
            None => Ok(None),
        }
    }

    fn delete_context(&mut self) -> Result<()> {
        self.conn
            .execute("DELETE FROM context WHERE id = 1", [])
            .map_err(StorageError::from)?;
        Ok(())
    }

    // ==================== Buffer Operations ====================

    #[allow(clippy::cast_possible_wrap)]
    fn add_buffer(&mut self, buffer: &Buffer) -> Result<i64> {
        let now = Self::now();

        self.conn
            .execute(
                r"
            INSERT INTO buffers (
                name, source_path, content, content_type, content_hash,
                size, line_count, chunk_count, created_at, updated_at
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        ",
                params![
                    buffer.name,
                    buffer
                        .source
                        .as_ref()
                        .map(|p| p.to_string_lossy().to_string()),
                    buffer.content,
                    buffer.metadata.content_type,
                    buffer.metadata.content_hash,
                    buffer.metadata.size as i64,
                    buffer.metadata.line_count.map(|c| c as i64),
                    buffer.metadata.chunk_count.map(|c| c as i64),
                    now,
                    now,
                ],
            )
            .map_err(StorageError::from)?;

        Ok(self.conn.last_insert_rowid())
    }

    fn get_buffer(&self, id: i64) -> Result<Option<Buffer>> {
        let result = self
            .conn
            .query_row(
                r"
            SELECT id, name, source_path, content, content_type, content_hash,
                   size, line_count, chunk_count, created_at, updated_at
            FROM buffers WHERE id = ?
        ",
                params![id],
                |row| {
                    Ok(Buffer {
                        id: Some(row.get::<_, i64>(0)?),
                        name: row.get(1)?,
                        source: row.get::<_, Option<String>>(2)?.map(PathBuf::from),
                        content: row.get(3)?,
                        metadata: BufferMetadata {
                            content_type: row.get(4)?,
                            content_hash: row.get(5)?,
                            size: row.get::<_, i64>(6)? as usize,
                            line_count: row.get::<_, Option<i64>>(7)?.map(|c| c as usize),
                            chunk_count: row.get::<_, Option<i64>>(8)?.map(|c| c as usize),
                            created_at: row.get(9)?,
                            updated_at: row.get(10)?,
                        },
                    })
                },
            )
            .optional()
            .map_err(StorageError::from)?;

        Ok(result)
    }

    fn get_buffer_by_name(&self, name: &str) -> Result<Option<Buffer>> {
        let id: Option<i64> = self
            .conn
            .query_row(
                "SELECT id FROM buffers WHERE name = ?",
                params![name],
                |row| row.get(0),
            )
            .optional()
            .map_err(StorageError::from)?;

        id.map_or(Ok(None), |id| self.get_buffer(id))
    }

    fn list_buffers(&self) -> Result<Vec<Buffer>> {
        let mut stmt = self
            .conn
            .prepare(
                r"
            SELECT id, name, source_path, content, content_type, content_hash,
                   size, line_count, chunk_count, created_at, updated_at
            FROM buffers ORDER BY id
        ",
            )
            .map_err(StorageError::from)?;

        let buffers = stmt
            .query_map([], |row| {
                Ok(Buffer {
                    id: Some(row.get::<_, i64>(0)?),
                    name: row.get(1)?,
                    source: row.get::<_, Option<String>>(2)?.map(PathBuf::from),
                    content: row.get(3)?,
                    metadata: BufferMetadata {
                        content_type: row.get(4)?,
                        content_hash: row.get(5)?,
                        size: row.get::<_, i64>(6)? as usize,
                        line_count: row.get::<_, Option<i64>>(7)?.map(|c| c as usize),
                        chunk_count: row.get::<_, Option<i64>>(8)?.map(|c| c as usize),
                        created_at: row.get(9)?,
                        updated_at: row.get(10)?,
                    },
                })
            })
            .map_err(StorageError::from)?
            .collect::<std::result::Result<Vec<_>, _>>()
            .map_err(StorageError::from)?;

        Ok(buffers)
    }

    #[allow(clippy::cast_possible_wrap)]
    fn update_buffer(&mut self, buffer: &Buffer) -> Result<()> {
        let id = buffer.id.ok_or_else(|| StorageError::BufferNotFound {
            identifier: "no ID".to_string(),
        })?;

        let now = Self::now();

        self.conn
            .execute(
                r"
            UPDATE buffers SET
                name = ?, source_path = ?, content = ?, content_type = ?,
                content_hash = ?, size = ?, line_count = ?, chunk_count = ?,
                updated_at = ?
            WHERE id = ?
        ",
                params![
                    buffer.name,
                    buffer
                        .source
                        .as_ref()
                        .map(|p| p.to_string_lossy().to_string()),
                    buffer.content,
                    buffer.metadata.content_type,
                    buffer.metadata.content_hash,
                    buffer.metadata.size as i64,
                    buffer.metadata.line_count.map(|c| c as i64),
                    buffer.metadata.chunk_count.map(|c| c as i64),
                    now,
                    id,
                ],
            )
            .map_err(StorageError::from)?;

        Ok(())
    }

    fn delete_buffer(&mut self, id: i64) -> Result<()> {
        // Chunks are deleted automatically via CASCADE
        self.conn
            .execute("DELETE FROM buffers WHERE id = ?", params![id])
            .map_err(StorageError::from)?;
        Ok(())
    }

    fn buffer_count(&self) -> Result<usize> {
        let count: i64 = self
            .conn
            .query_row("SELECT COUNT(*) FROM buffers", [], |row| row.get(0))
            .map_err(StorageError::from)?;
        Ok(count as usize)
    }

    // ==================== Chunk Operations ====================

    #[allow(clippy::cast_possible_wrap)]
    fn add_chunks(&mut self, buffer_id: i64, chunks: &[Chunk]) -> Result<()> {
        let tx = self.conn.transaction().map_err(StorageError::from)?;
        let now = Self::now();

        {
            let mut stmt = tx
                .prepare(
                    r"
                INSERT INTO chunks (
                    buffer_id, content, byte_start, byte_end, chunk_index,
                    strategy, token_count, line_start, line_end, has_overlap,
                    content_hash, custom_metadata, created_at
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            ",
                )
                .map_err(StorageError::from)?;

            for chunk in chunks {
                let custom_meta = chunk.metadata.custom.clone();

                let (line_start, line_end) = chunk
                    .metadata
                    .line_range
                    .as_ref()
                    .map_or((None, None), |r| (Some(r.start as i64), Some(r.end as i64)));

                stmt.execute(params![
                    buffer_id,
                    chunk.content,
                    chunk.byte_range.start as i64,
                    chunk.byte_range.end as i64,
                    chunk.index as i64,
                    chunk.metadata.strategy,
                    chunk.metadata.token_count.map(|c| c as i64),
                    line_start,
                    line_end,
                    i64::from(chunk.metadata.has_overlap),
                    chunk.metadata.content_hash,
                    custom_meta,
                    now,
                ])
                .map_err(StorageError::from)?;
            }
        }

        tx.commit().map_err(StorageError::from)?;

        // Update chunk count on buffer
        self.conn
            .execute(
                "UPDATE buffers SET chunk_count = ? WHERE id = ?",
                params![chunks.len() as i64, buffer_id],
            )
            .map_err(StorageError::from)?;

        Ok(())
    }

    fn get_chunks(&self, buffer_id: i64) -> Result<Vec<Chunk>> {
        let mut stmt = self
            .conn
            .prepare(
                r"
            SELECT id, buffer_id, content, byte_start, byte_end, chunk_index,
                   strategy, token_count, line_start, line_end, has_overlap,
                   content_hash, custom_metadata, created_at
            FROM chunks WHERE buffer_id = ? ORDER BY chunk_index
        ",
            )
            .map_err(StorageError::from)?;

        let chunks = stmt
            .query_map(params![buffer_id], |row| {
                let line_start: Option<i64> = row.get(8)?;
                let line_end: Option<i64> = row.get(9)?;
                let line_range = match (line_start, line_end) {
                    (Some(s), Some(e)) => Some((s as usize)..(e as usize)),
                    _ => None,
                };

                Ok(Chunk {
                    id: Some(row.get::<_, i64>(0)?),
                    buffer_id: row.get(1)?,
                    content: row.get(2)?,
                    byte_range: (row.get::<_, i64>(3)? as usize)..(row.get::<_, i64>(4)? as usize),
                    index: row.get::<_, i64>(5)? as usize,
                    metadata: ChunkMetadata {
                        strategy: row.get(6)?,
                        token_count: row.get::<_, Option<i64>>(7)?.map(|c| c as usize),
                        line_range,
                        has_overlap: row.get::<_, i64>(10)? != 0,
                        content_hash: row.get(11)?,
                        custom: row.get(12)?,
                        created_at: row.get(13)?,
                    },
                })
            })
            .map_err(StorageError::from)?
            .collect::<std::result::Result<Vec<_>, _>>()
            .map_err(StorageError::from)?;

        Ok(chunks)
    }

    fn get_chunk(&self, id: i64) -> Result<Option<Chunk>> {
        let result = self
            .conn
            .query_row(
                r"
            SELECT id, buffer_id, content, byte_start, byte_end, chunk_index,
                   strategy, token_count, line_start, line_end, has_overlap,
                   content_hash, custom_metadata, created_at
            FROM chunks WHERE id = ?
        ",
                params![id],
                |row| {
                    let line_start: Option<i64> = row.get(8)?;
                    let line_end: Option<i64> = row.get(9)?;
                    let line_range = match (line_start, line_end) {
                        (Some(s), Some(e)) => Some((s as usize)..(e as usize)),
                        _ => None,
                    };

                    Ok(Chunk {
                        id: Some(row.get::<_, i64>(0)?),
                        buffer_id: row.get(1)?,
                        content: row.get(2)?,
                        byte_range: (row.get::<_, i64>(3)? as usize)
                            ..(row.get::<_, i64>(4)? as usize),
                        index: row.get::<_, i64>(5)? as usize,
                        metadata: ChunkMetadata {
                            strategy: row.get(6)?,
                            token_count: row.get::<_, Option<i64>>(7)?.map(|c| c as usize),
                            line_range,
                            has_overlap: row.get::<_, i64>(10)? != 0,
                            content_hash: row.get(11)?,
                            custom: row.get(12)?,
                            created_at: row.get(13)?,
                        },
                    })
                },
            )
            .optional()
            .map_err(StorageError::from)?;

        Ok(result)
    }

    fn delete_chunks(&mut self, buffer_id: i64) -> Result<()> {
        self.conn
            .execute("DELETE FROM chunks WHERE buffer_id = ?", params![buffer_id])
            .map_err(StorageError::from)?;

        // Update chunk count on buffer
        self.conn
            .execute(
                "UPDATE buffers SET chunk_count = 0 WHERE id = ?",
                params![buffer_id],
            )
            .map_err(StorageError::from)?;

        Ok(())
    }

    fn chunk_count(&self, buffer_id: i64) -> Result<usize> {
        let count: i64 = self
            .conn
            .query_row(
                "SELECT COUNT(*) FROM chunks WHERE buffer_id = ?",
                params![buffer_id],
                |row| row.get(0),
            )
            .map_err(StorageError::from)?;
        Ok(count as usize)
    }

    // ==================== Utility Operations ====================

    fn export_buffers(&self) -> Result<String> {
        let buffers = self.list_buffers()?;
        let mut output = String::new();

        for (i, buffer) in buffers.iter().enumerate() {
            if i > 0 {
                output.push_str("\n\n");
            }
            output.push_str(&buffer.content);
        }

        Ok(output)
    }

    fn stats(&self) -> Result<StorageStats> {
        let buffer_count = self.buffer_count()?;

        let chunk_count: i64 = self
            .conn
            .query_row("SELECT COUNT(*) FROM chunks", [], |row| row.get(0))
            .map_err(StorageError::from)?;

        let total_size: i64 = self
            .conn
            .query_row("SELECT COALESCE(SUM(size), 0) FROM buffers", [], |row| {
                row.get(0)
            })
            .map_err(StorageError::from)?;

        let has_context = self.load_context()?.is_some();

        let schema_version = self.get_schema_version()?.unwrap_or(0);

        let db_size = self
            .path
            .as_ref()
            .and_then(|p| std::fs::metadata(p).ok().map(|m| m.len()));

        Ok(StorageStats {
            buffer_count,
            chunk_count: chunk_count as usize,
            total_content_size: total_size as usize,
            has_context,
            schema_version,
            db_size,
        })
    }
}

// ==================== Embedding & Search Operations ====================

impl SqliteStorage {
    /// Retrieves multiple chunks by their IDs in a single query.
    ///
    /// More efficient than calling [`Storage::get_chunk`] repeatedly when fetching
    /// several chunks at once (e.g., populating search result previews).
    ///
    /// Returns a map from chunk ID to [`Chunk`].
    ///
    /// # Errors
    ///
    /// Returns an error if the query fails.
    pub fn get_chunks_by_ids(&self, ids: &[i64]) -> Result<std::collections::HashMap<i64, Chunk>> {
        if ids.is_empty() {
            return Ok(std::collections::HashMap::new());
        }

        let placeholders = ids.iter().map(|_| "?").collect::<Vec<_>>().join(",");
        let sql = format!(
            "SELECT id, buffer_id, content, byte_start, byte_end, chunk_index, \
             strategy, token_count, line_start, line_end, has_overlap, \
             content_hash, custom_metadata, created_at \
             FROM chunks WHERE id IN ({placeholders})"
        );

        let mut stmt = self.conn.prepare(&sql).map_err(StorageError::from)?;

        let chunks = stmt
            .query_map(rusqlite::params_from_iter(ids.iter().copied()), |row| {
                let line_start: Option<i64> = row.get(8)?;
                let line_end: Option<i64> = row.get(9)?;
                let line_range = match (line_start, line_end) {
                    (Some(s), Some(e)) => Some((s as usize)..(e as usize)),
                    _ => None,
                };

                Ok(Chunk {
                    id: Some(row.get::<_, i64>(0)?),
                    buffer_id: row.get(1)?,
                    content: row.get(2)?,
                    byte_range: (row.get::<_, i64>(3)? as usize)..(row.get::<_, i64>(4)? as usize),
                    index: row.get::<_, i64>(5)? as usize,
                    metadata: ChunkMetadata {
                        strategy: row.get(6)?,
                        token_count: row.get::<_, Option<i64>>(7)?.map(|c| c as usize),
                        line_range,
                        has_overlap: row.get::<_, i64>(10)? != 0,
                        content_hash: row.get(11)?,
                        custom: row.get(12)?,
                        created_at: row.get(13)?,
                    },
                })
            })
            .map_err(StorageError::from)?
            .collect::<std::result::Result<Vec<_>, _>>()
            .map_err(StorageError::from)?;

        Ok(chunks
            .into_iter()
            .filter_map(|c| c.id.map(|id| (id, c)))
            .collect())
    }

    /// Stores an embedding for a chunk.
    ///
    /// # Arguments
    ///
    /// * `chunk_id` - The chunk ID to associate the embedding with.
    /// * `embedding` - The embedding vector (f32 array).
    /// * `model_name` - Optional name of the model that generated the embedding.
    ///
    /// # Errors
    ///
    /// Returns an error if the embedding cannot be stored.
    #[allow(clippy::cast_possible_wrap)]
    pub fn store_embedding(
        &mut self,
        chunk_id: i64,
        embedding: &[f32],
        model_name: Option<&str>,
    ) -> Result<()> {
        let now = Self::now();

        // Serialize f32 array to bytes (little-endian)
        let mut bytes = Vec::with_capacity(embedding.len() * 4);
        for f in embedding {
            bytes.extend_from_slice(&f.to_le_bytes());
        }

        self.conn
            .execute(
                r"
                INSERT OR REPLACE INTO chunk_embeddings (chunk_id, embedding, dimensions, model_name, created_at)
                VALUES (?, ?, ?, ?, ?)
            ",
                params![chunk_id, bytes, embedding.len() as i64, model_name, now],
            )
            .map_err(StorageError::from)?;

        Ok(())
    }

    /// Retrieves the embedding for a chunk.
    ///
    /// # Errors
    ///
    /// Returns an error if the query fails.
    pub fn get_embedding(&self, chunk_id: i64) -> Result<Option<Vec<f32>>> {
        let result: Option<Vec<u8>> = self
            .conn
            .query_row(
                "SELECT embedding FROM chunk_embeddings WHERE chunk_id = ?",
                params![chunk_id],
                |row| row.get(0),
            )
            .optional()
            .map_err(StorageError::from)?;

        Ok(result.map(|bytes| {
            bytes
                .chunks_exact(4)
                .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
                .collect()
        }))
    }

    /// Gets the distinct model names used for embeddings in a buffer.
    ///
    /// Returns the set of model names used to generate embeddings for
    /// chunks belonging to the specified buffer.
    ///
    /// # Errors
    ///
    /// Returns an error if the query fails.
    pub fn get_embedding_models(&self, buffer_id: i64) -> Result<Vec<String>> {
        let mut stmt = self
            .conn
            .prepare(
                r"
                SELECT DISTINCT ce.model_name
                FROM chunk_embeddings ce
                JOIN chunks c ON ce.chunk_id = c.id
                WHERE c.buffer_id = ? AND ce.model_name IS NOT NULL
                ",
            )
            .map_err(StorageError::from)?;

        let models = stmt
            .query_map(params![buffer_id], |row| row.get::<_, String>(0))
            .map_err(StorageError::from)?
            .filter_map(std::result::Result::ok)
            .collect();

        Ok(models)
    }

    /// Gets the count of embeddings by model name for a buffer.
    ///
    /// Returns a list of (`model_name`, count) pairs.
    ///
    /// # Errors
    ///
    /// Returns an error if the query fails.
    pub fn get_embedding_model_counts(&self, buffer_id: i64) -> Result<Vec<(Option<String>, i64)>> {
        let mut stmt = self
            .conn
            .prepare(
                r"
                SELECT ce.model_name, COUNT(*) as count
                FROM chunk_embeddings ce
                JOIN chunks c ON ce.chunk_id = c.id
                WHERE c.buffer_id = ?
                GROUP BY ce.model_name
                ",
            )
            .map_err(StorageError::from)?;

        let counts = stmt
            .query_map(params![buffer_id], |row| {
                Ok((row.get::<_, Option<String>>(0)?, row.get::<_, i64>(1)?))
            })
            .map_err(StorageError::from)?
            .filter_map(std::result::Result::ok)
            .collect();

        Ok(counts)
    }

    /// Stores embeddings for multiple chunks in a batch.
    ///
    /// # Errors
    ///
    /// Returns an error if any embedding cannot be stored.
    #[allow(clippy::cast_possible_wrap)]
    pub fn store_embeddings_batch(
        &mut self,
        embeddings: &[(i64, Vec<f32>)],
        model_name: Option<&str>,
    ) -> Result<()> {
        let tx = self.conn.transaction().map_err(StorageError::from)?;
        let now = Self::now();

        {
            let mut stmt = tx
                .prepare(
                    r"
                    INSERT OR REPLACE INTO chunk_embeddings (chunk_id, embedding, dimensions, model_name, created_at)
                    VALUES (?, ?, ?, ?, ?)
                ",
                )
                .map_err(StorageError::from)?;

            for (chunk_id, embedding) in embeddings {
                let mut bytes = Vec::with_capacity(embedding.len() * 4);
                bytes.extend(embedding.iter().flat_map(|f| f.to_le_bytes()));

                stmt.execute(params![
                    chunk_id,
                    bytes,
                    embedding.len() as i64,
                    model_name,
                    now
                ])
                .map_err(StorageError::from)?;
            }
        }

        tx.commit().map_err(StorageError::from)?;
        Ok(())
    }

    /// Deletes the embedding for a chunk.
    ///
    /// # Errors
    ///
    /// Returns an error if deletion fails.
    pub fn delete_embedding(&mut self, chunk_id: i64) -> Result<()> {
        self.conn
            .execute(
                "DELETE FROM chunk_embeddings WHERE chunk_id = ?",
                params![chunk_id],
            )
            .map_err(StorageError::from)?;
        Ok(())
    }

    /// Performs FTS5 BM25 full-text search.
    ///
    /// Returns chunk IDs and their BM25 scores (higher is better match).
    ///
    /// # Arguments
    ///
    /// * `query` - The search query (supports FTS5 query syntax).
    /// * `limit` - Maximum number of results to return.
    ///
    /// # Errors
    ///
    /// Returns an error if the search fails.
    #[allow(clippy::cast_possible_wrap)]
    pub fn search_fts(&self, query: &str, limit: usize) -> Result<Vec<(i64, f64)>> {
        self.search_fts_in_buffer(query, limit, None)
    }

    /// Performs FTS5 BM25 full-text search, optionally restricted to a buffer.
    ///
    /// # Errors
    ///
    /// Returns an error if the search fails.
    #[allow(clippy::cast_possible_wrap)]
    pub fn search_fts_in_buffer(
        &self,
        query: &str,
        limit: usize,
        buffer_id: Option<i64>,
    ) -> Result<Vec<(i64, f64)>> {
        // FTS5 bm25() returns negative scores, more negative = better match
        // We negate it so higher scores = better match

        // Convert space-separated terms to OR query for more forgiving search
        // Each term is quoted to escape FTS5 special characters (?, *, ^, etc.)
        // "CLI tool?" becomes '"CLI" OR "tool?"' so special chars are treated as literals
        let fts_query = query
            .split_whitespace()
            .map(|term| format!("\"{}\"", term.replace('"', "\"\"")))
            .collect::<Vec<_>>()
            .join(" OR ");

        let results = if let Some(buffer_id) = buffer_id {
            let mut stmt = self
                .conn
                .prepare(
                    r"
                    SELECT chunks_fts.rowid, -bm25(chunks_fts) as score
                    FROM chunks_fts
                    INNER JOIN chunks ON chunks.id = chunks_fts.rowid
                    WHERE chunks_fts MATCH ?1
                      AND chunks.buffer_id = ?2
                    ORDER BY score DESC
                    LIMIT ?3
                ",
                )
                .map_err(StorageError::from)?;

            stmt.query_map(params![fts_query, buffer_id, limit as i64], |row| {
                Ok((row.get::<_, i64>(0)?, row.get::<_, f64>(1)?))
            })
            .map_err(StorageError::from)?
            .collect::<std::result::Result<Vec<_>, _>>()
            .map_err(StorageError::from)?
        } else {
            let mut stmt = self
                .conn
                .prepare(
                    r"
                    SELECT rowid, -bm25(chunks_fts) as score
                    FROM chunks_fts
                    WHERE chunks_fts MATCH ?1
                    ORDER BY score DESC
                    LIMIT ?2
                ",
                )
                .map_err(StorageError::from)?;

            stmt.query_map(params![fts_query, limit as i64], |row| {
                Ok((row.get::<_, i64>(0)?, row.get::<_, f64>(1)?))
            })
            .map_err(StorageError::from)?
            .collect::<std::result::Result<Vec<_>, _>>()
            .map_err(StorageError::from)?
        };

        Ok(results)
    }

    /// Returns all chunk embeddings for vector similarity search.
    ///
    /// # Errors
    ///
    /// Returns an error if the query fails.
    pub fn get_all_embeddings(&self) -> Result<Vec<(i64, Vec<f32>)>> {
        let mut stmt = self
            .conn
            .prepare("SELECT chunk_id, embedding FROM chunk_embeddings")
            .map_err(StorageError::from)?;

        let results = stmt
            .query_map([], |row| {
                let chunk_id: i64 = row.get(0)?;
                let bytes: Vec<u8> = row.get(1)?;
                let embedding: Vec<f32> = bytes
                    .chunks_exact(4)
                    .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
                    .collect();
                Ok((chunk_id, embedding))
            })
            .map_err(StorageError::from)?
            .collect::<std::result::Result<Vec<_>, _>>()
            .map_err(StorageError::from)?;

        Ok(results)
    }

    /// Returns chunk embeddings belonging to one buffer for vector similarity search.
    ///
    /// # Errors
    ///
    /// Returns an error if the query fails.
    pub fn get_embeddings_in_buffer(&self, buffer_id: i64) -> Result<Vec<(i64, Vec<f32>)>> {
        let mut stmt = self
            .conn
            .prepare(
                r"
                SELECT ce.chunk_id, ce.embedding
                FROM chunk_embeddings ce
                INNER JOIN chunks c ON c.id = ce.chunk_id
                WHERE c.buffer_id = ?
            ",
            )
            .map_err(StorageError::from)?;

        let results = stmt
            .query_map(params![buffer_id], |row| {
                let chunk_id: i64 = row.get(0)?;
                let bytes: Vec<u8> = row.get(1)?;
                let embedding: Vec<f32> = bytes
                    .chunks_exact(4)
                    .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
                    .collect();
                Ok((chunk_id, embedding))
            })
            .map_err(StorageError::from)?
            .collect::<std::result::Result<Vec<_>, _>>()
            .map_err(StorageError::from)?;

        Ok(results)
    }

    /// Counts chunks with embeddings.
    ///
    /// # Errors
    ///
    /// Returns an error if the count fails.
    pub fn embedding_count(&self) -> Result<usize> {
        let count: i64 = self
            .conn
            .query_row("SELECT COUNT(*) FROM chunk_embeddings", [], |row| {
                row.get(0)
            })
            .map_err(StorageError::from)?;
        Ok(count as usize)
    }

    /// Returns `true` if every chunk in the buffer has an embedding (or the buffer has no chunks).
    ///
    /// Uses a single `NOT EXISTS` query instead of per-chunk lookups, making it O(1) in terms
    /// of round-trips regardless of how many chunks the buffer contains.
    ///
    /// # Errors
    ///
    /// Returns an error if the query fails.
    pub fn all_chunks_have_embeddings(&self, buffer_id: i64) -> Result<bool> {
        let result: i64 = self
            .conn
            .query_row(
                r"
                SELECT NOT EXISTS (
                    SELECT 1 FROM chunks c
                    LEFT JOIN chunk_embeddings e ON e.chunk_id = c.id
                    WHERE c.buffer_id = ? AND e.chunk_id IS NULL
                )
                ",
                params![buffer_id],
                |row| row.get(0),
            )
            .map_err(StorageError::from)?;
        Ok(result != 0)
    }

    /// Checks if a chunk has an embedding.
    ///
    /// # Errors
    ///
    /// Returns an error if the query fails.
    pub fn has_embedding(&self, chunk_id: i64) -> Result<bool> {
        let count: i64 = self
            .conn
            .query_row(
                "SELECT COUNT(*) FROM chunk_embeddings WHERE chunk_id = ?",
                params![chunk_id],
                |row| row.get(0),
            )
            .map_err(StorageError::from)?;
        Ok(count > 0)
    }

    /// Gets chunk IDs that need embedding (either no embedding or wrong model).
    ///
    /// This is used for incremental embedding updates. Returns chunks that:
    /// - Have no embedding at all, OR
    /// - Have an embedding with a different model name (if `current_model` is provided)
    ///
    /// # Arguments
    ///
    /// * `buffer_id` - The buffer to check.
    /// * `current_model` - Optional model name to check against. If provided,
    ///   chunks with different models are included.
    ///
    /// # Errors
    ///
    /// Returns an error if the query fails.
    pub fn get_chunks_needing_embedding(
        &self,
        buffer_id: i64,
        current_model: Option<&str>,
    ) -> Result<Vec<i64>> {
        let mut results = Vec::new();

        // Get chunks without any embedding
        let mut stmt = self
            .conn
            .prepare(
                r"
                SELECT c.id FROM chunks c
                LEFT JOIN chunk_embeddings e ON c.id = e.chunk_id
                WHERE c.buffer_id = ? AND e.chunk_id IS NULL
                ",
            )
            .map_err(StorageError::from)?;

        let rows = stmt
            .query_map(params![buffer_id], |row| row.get(0))
            .map_err(StorageError::from)?;

        for row in rows {
            results.push(row.map_err(StorageError::from)?);
        }

        // If model specified, also get chunks with different model
        if let Some(model) = current_model {
            let mut stmt = self
                .conn
                .prepare(
                    r"
                    SELECT c.id FROM chunks c
                    INNER JOIN chunk_embeddings e ON c.id = e.chunk_id
                    WHERE c.buffer_id = ? AND (e.model_name IS NULL OR e.model_name != ?)
                    ",
                )
                .map_err(StorageError::from)?;

            let rows = stmt
                .query_map(params![buffer_id, model], |row| row.get(0))
                .map_err(StorageError::from)?;

            for row in rows {
                results.push(row.map_err(StorageError::from)?);
            }
        }

        // Deduplicate (in case of overlap, though shouldn't happen)
        results.sort_unstable();
        results.dedup();
        Ok(results)
    }

    /// Gets chunks without any embedding for a buffer.
    ///
    /// Simpler version of `get_chunks_needing_embedding` when model doesn't matter.
    ///
    /// # Errors
    ///
    /// Returns an error if the query fails.
    pub fn get_chunks_without_embedding(&self, buffer_id: i64) -> Result<Vec<i64>> {
        self.get_chunks_needing_embedding(buffer_id, None)
    }

    /// Deletes embeddings with a specific model name.
    ///
    /// Useful for cleaning up embeddings from old models before re-embedding.
    ///
    /// # Arguments
    ///
    /// * `buffer_id` - The buffer to clean.
    /// * `model_name` - The model name to match (or None to match NULL).
    ///
    /// # Returns
    ///
    /// The number of embeddings deleted.
    ///
    /// # Errors
    ///
    /// Returns an error if deletion fails.
    pub fn delete_embeddings_by_model(
        &mut self,
        buffer_id: i64,
        model_name: Option<&str>,
    ) -> Result<usize> {
        let deleted = match model_name {
            Some(name) => self
                .conn
                .execute(
                    r"
                    DELETE FROM chunk_embeddings
                    WHERE chunk_id IN (
                        SELECT id FROM chunks WHERE buffer_id = ?
                    ) AND model_name = ?
                    ",
                    params![buffer_id, name],
                )
                .map_err(StorageError::from)?,
            None => self
                .conn
                .execute(
                    r"
                    DELETE FROM chunk_embeddings
                    WHERE chunk_id IN (
                        SELECT id FROM chunks WHERE buffer_id = ?
                    ) AND model_name IS NULL
                    ",
                    params![buffer_id],
                )
                .map_err(StorageError::from)?,
        };
        Ok(deleted)
    }

    /// Gets embedding statistics for a buffer.
    ///
    /// Returns counts of embedded vs total chunks, and model breakdown.
    ///
    /// # Errors
    ///
    /// Returns an error if the query fails.
    pub fn get_embedding_stats(&self, buffer_id: i64) -> Result<EmbeddingStats> {
        // Total chunks
        let total_chunks: i64 = self
            .conn
            .query_row(
                "SELECT COUNT(*) FROM chunks WHERE buffer_id = ?",
                params![buffer_id],
                |row| row.get(0),
            )
            .map_err(StorageError::from)?;

        // Embedded chunks
        let embedded_chunks: i64 = self
            .conn
            .query_row(
                r"
                SELECT COUNT(*) FROM chunk_embeddings e
                INNER JOIN chunks c ON e.chunk_id = c.id
                WHERE c.buffer_id = ?
                ",
                params![buffer_id],
                |row| row.get(0),
            )
            .map_err(StorageError::from)?;

        // Model counts
        let model_counts = self.get_embedding_model_counts(buffer_id)?;

        Ok(EmbeddingStats {
            total_chunks: total_chunks as usize,
            embedded_chunks: embedded_chunks as usize,
            model_counts,
        })
    }
}

/// Statistics about embeddings for a buffer.
#[derive(Debug, Clone)]
pub struct EmbeddingStats {
    /// Total number of chunks in the buffer.
    pub total_chunks: usize,
    /// Number of chunks with embeddings.
    pub embedded_chunks: usize,
    /// Count of embeddings by model (`model_name`, count).
    pub model_counts: Vec<(Option<String>, i64)>,
}

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

    fn setup() -> SqliteStorage {
        let mut storage = SqliteStorage::in_memory().unwrap();
        storage.init().unwrap();
        storage
    }

    #[test]
    fn test_init() {
        let mut storage = SqliteStorage::in_memory().unwrap();
        assert!(storage.init().is_ok());
        assert!(storage.is_initialized().unwrap());
    }

    #[test]
    fn test_init_idempotent() {
        let mut storage = SqliteStorage::in_memory().unwrap();
        assert!(storage.init().is_ok());
        assert!(storage.init().is_ok()); // Second init should be fine
    }

    #[test]
    fn test_context_crud() {
        let mut storage = setup();

        // No context initially
        assert!(storage.load_context().unwrap().is_none());

        // Save context
        let mut ctx = Context::new();
        ctx.set_variable("key".to_string(), ContextValue::String("value".to_string()));
        storage.save_context(&ctx).unwrap();

        // Load context
        let loaded = storage.load_context().unwrap().unwrap();
        assert_eq!(
            loaded.get_variable("key"),
            Some(&ContextValue::String("value".to_string()))
        );

        // Delete context
        storage.delete_context().unwrap();
        assert!(storage.load_context().unwrap().is_none());
    }

    #[test]
    fn test_buffer_crud() {
        let mut storage = setup();

        // Add buffer
        let buffer = Buffer::from_named("test".to_string(), "Hello, world!".to_string());
        let id = storage.add_buffer(&buffer).unwrap();
        assert!(id > 0);

        // Get buffer
        let loaded = storage.get_buffer(id).unwrap().unwrap();
        assert_eq!(loaded.name, Some("test".to_string()));
        assert_eq!(loaded.content, "Hello, world!");

        // Get by name
        let by_name = storage.get_buffer_by_name("test").unwrap().unwrap();
        assert_eq!(by_name.id, Some(id));

        // List buffers
        let buffers = storage.list_buffers().unwrap();
        assert_eq!(buffers.len(), 1);

        // Update buffer
        let mut updated = loaded;
        updated.content = "Updated content".to_string();
        storage.update_buffer(&updated).unwrap();

        let reloaded = storage.get_buffer(id).unwrap().unwrap();
        assert_eq!(reloaded.content, "Updated content");

        // Delete buffer
        storage.delete_buffer(id).unwrap();
        assert!(storage.get_buffer(id).unwrap().is_none());
    }

    #[test]
    fn test_chunk_crud() {
        let mut storage = setup();

        // Create buffer first
        let buffer = Buffer::from_content("Hello, world!".to_string());
        let buffer_id = storage.add_buffer(&buffer).unwrap();

        // Add chunks
        let chunks = vec![
            Chunk::new(buffer_id, "Hello, ".to_string(), 0..7, 0),
            Chunk::new(buffer_id, "world!".to_string(), 7..13, 1),
        ];
        storage.add_chunks(buffer_id, &chunks).unwrap();

        // Get chunks
        let loaded = storage.get_chunks(buffer_id).unwrap();
        assert_eq!(loaded.len(), 2);
        assert_eq!(loaded[0].content, "Hello, ");
        assert_eq!(loaded[1].content, "world!");

        // Chunk count
        assert_eq!(storage.chunk_count(buffer_id).unwrap(), 2);

        // Get single chunk
        let chunk_id = loaded[0].id.unwrap();
        let single = storage.get_chunk(chunk_id).unwrap().unwrap();
        assert_eq!(single.content, "Hello, ");

        // Delete chunks
        storage.delete_chunks(buffer_id).unwrap();
        assert_eq!(storage.chunk_count(buffer_id).unwrap(), 0);
    }

    #[test]
    fn test_cascade_delete() {
        let mut storage = setup();

        // Create buffer with chunks
        let buffer = Buffer::from_content("Hello, world!".to_string());
        let buffer_id = storage.add_buffer(&buffer).unwrap();

        let chunks = vec![Chunk::new(buffer_id, "Hello".to_string(), 0..5, 0)];
        storage.add_chunks(buffer_id, &chunks).unwrap();

        // Verify chunk exists
        assert_eq!(storage.chunk_count(buffer_id).unwrap(), 1);

        // Delete buffer - chunks should be deleted too
        storage.delete_buffer(buffer_id).unwrap();

        // Verify no orphan chunks (query all chunks)
        let count: i64 = storage
            .conn
            .query_row("SELECT COUNT(*) FROM chunks", [], |row| row.get(0))
            .unwrap();
        assert_eq!(count, 0);
    }

    #[test]
    fn test_reset() {
        let mut storage = setup();

        // Add some data
        let ctx = Context::new();
        storage.save_context(&ctx).unwrap();

        let buffer = Buffer::from_content("test".to_string());
        storage.add_buffer(&buffer).unwrap();

        // Reset
        storage.reset().unwrap();

        // Verify empty
        assert!(storage.load_context().unwrap().is_none());
        assert_eq!(storage.buffer_count().unwrap(), 0);
    }

    #[test]
    fn test_stats() {
        let mut storage = setup();

        // Empty stats
        let stats = storage.stats().unwrap();
        assert_eq!(stats.buffer_count, 0);
        assert_eq!(stats.chunk_count, 0);
        assert!(!stats.has_context);

        // Add data
        let ctx = Context::new();
        storage.save_context(&ctx).unwrap();

        let buffer = Buffer::from_content("Hello, world!".to_string());
        let buffer_id = storage.add_buffer(&buffer).unwrap();

        let chunks = vec![Chunk::new(buffer_id, "Hello".to_string(), 0..5, 0)];
        storage.add_chunks(buffer_id, &chunks).unwrap();

        // Stats with data
        let stats = storage.stats().unwrap();
        assert_eq!(stats.buffer_count, 1);
        assert_eq!(stats.chunk_count, 1);
        assert!(stats.has_context);
        assert_eq!(stats.total_content_size, 13);
    }

    #[test]
    fn test_export_buffers() {
        let mut storage = setup();

        storage
            .add_buffer(&Buffer::from_content("First".to_string()))
            .unwrap();
        storage
            .add_buffer(&Buffer::from_content("Second".to_string()))
            .unwrap();

        let exported = storage.export_buffers().unwrap();
        assert_eq!(exported, "First\n\nSecond");
    }

    // Helper: create a buffer with one chunk and return (buffer_id, chunk_id).
    fn setup_buffer_with_chunk(storage: &mut SqliteStorage) -> (i64, i64) {
        let buffer = Buffer::from_content("test content".to_string());
        let buffer_id = storage.add_buffer(&buffer).unwrap();
        let chunks = vec![Chunk::new(buffer_id, "test content".to_string(), 0..12, 0)];
        storage.add_chunks(buffer_id, &chunks).unwrap();
        let chunk_id = storage.get_chunks(buffer_id).unwrap()[0].id.unwrap();
        (buffer_id, chunk_id)
    }

    #[test]
    fn test_store_and_get_embedding() {
        let mut storage = setup();
        let (_buffer_id, chunk_id) = setup_buffer_with_chunk(&mut storage);

        let embedding = vec![0.1_f32, 0.2, 0.3, 0.4];
        storage
            .store_embedding(chunk_id, &embedding, Some("test-model"))
            .unwrap();

        let loaded = storage.get_embedding(chunk_id).unwrap().unwrap();
        assert_eq!(loaded.len(), 4);
        for (a, b) in loaded.iter().zip(embedding.iter()) {
            assert!((a - b).abs() < 1e-6, "expected {b}, got {a}");
        }
    }

    #[test]
    fn test_get_embedding_nonexistent() {
        let storage = setup();
        let result = storage.get_embedding(9999).unwrap();
        assert!(result.is_none());
    }

    #[test]
    fn test_store_embedding_upsert() {
        let mut storage = setup();
        let (_buffer_id, chunk_id) = setup_buffer_with_chunk(&mut storage);

        let embedding1 = vec![0.1_f32, 0.2];
        storage
            .store_embedding(chunk_id, &embedding1, Some("model-a"))
            .unwrap();

        // Upsert with new values
        let embedding2 = vec![0.9_f32, 0.8];
        storage
            .store_embedding(chunk_id, &embedding2, Some("model-a"))
            .unwrap();

        let loaded = storage.get_embedding(chunk_id).unwrap().unwrap();
        for (a, b) in loaded.iter().zip(embedding2.iter()) {
            assert!((a - b).abs() < 1e-6);
        }
    }

    #[test]
    fn test_has_embedding() {
        let mut storage = setup();
        let (_buffer_id, chunk_id) = setup_buffer_with_chunk(&mut storage);

        assert!(!storage.has_embedding(chunk_id).unwrap());

        storage.store_embedding(chunk_id, &[0.1_f32], None).unwrap();

        assert!(storage.has_embedding(chunk_id).unwrap());
    }

    #[test]
    fn test_embedding_count() {
        let mut storage = setup();
        assert_eq!(storage.embedding_count().unwrap(), 0);

        let (buffer_id, chunk_id) = setup_buffer_with_chunk(&mut storage);

        // Add a second chunk
        let chunks2 = vec![Chunk::new(buffer_id, "more".to_string(), 0..4, 1)];
        storage.add_chunks(buffer_id, &chunks2).unwrap();
        let chunk_id2 = storage
            .get_chunks(buffer_id)
            .unwrap()
            .into_iter()
            .find(|c| c.id != Some(chunk_id))
            .unwrap()
            .id
            .unwrap();

        storage.store_embedding(chunk_id, &[0.1_f32], None).unwrap();
        assert_eq!(storage.embedding_count().unwrap(), 1);

        storage
            .store_embedding(chunk_id2, &[0.2_f32], None)
            .unwrap();
        assert_eq!(storage.embedding_count().unwrap(), 2);
    }

    #[test]
    fn test_delete_embedding() {
        let mut storage = setup();
        let (_buffer_id, chunk_id) = setup_buffer_with_chunk(&mut storage);

        storage.store_embedding(chunk_id, &[0.1_f32], None).unwrap();
        assert!(storage.has_embedding(chunk_id).unwrap());

        storage.delete_embedding(chunk_id).unwrap();
        assert!(!storage.has_embedding(chunk_id).unwrap());
    }

    #[test]
    fn test_store_embeddings_batch() {
        let mut storage = setup();
        let (buffer_id, chunk_id1) = setup_buffer_with_chunk(&mut storage);

        // Add second chunk
        let chunks2 = vec![Chunk::new(buffer_id, "second".to_string(), 0..6, 1)];
        storage.add_chunks(buffer_id, &chunks2).unwrap();
        let chunk_id2 = storage
            .get_chunks(buffer_id)
            .unwrap()
            .into_iter()
            .find(|c| c.id != Some(chunk_id1))
            .unwrap()
            .id
            .unwrap();

        let batch = vec![
            (chunk_id1, vec![0.1_f32, 0.2]),
            (chunk_id2, vec![0.3_f32, 0.4]),
        ];
        storage
            .store_embeddings_batch(&batch, Some("batch-model"))
            .unwrap();

        assert!(storage.has_embedding(chunk_id1).unwrap());
        assert!(storage.has_embedding(chunk_id2).unwrap());
        assert_eq!(storage.embedding_count().unwrap(), 2);
    }

    #[test]
    fn test_get_all_embeddings() {
        let mut storage = setup();
        let (buffer_id, chunk_id1) = setup_buffer_with_chunk(&mut storage);

        let chunks2 = vec![Chunk::new(buffer_id, "second".to_string(), 0..6, 1)];
        storage.add_chunks(buffer_id, &chunks2).unwrap();
        let chunk_id2 = storage
            .get_chunks(buffer_id)
            .unwrap()
            .into_iter()
            .find(|c| c.id != Some(chunk_id1))
            .unwrap()
            .id
            .unwrap();

        storage
            .store_embedding(chunk_id1, &[0.1_f32], Some("m"))
            .unwrap();
        storage
            .store_embedding(chunk_id2, &[0.2_f32], Some("m"))
            .unwrap();
        let (other_buffer_id, other_chunk_id) = setup_buffer_with_chunk(&mut storage);
        storage
            .store_embedding(other_chunk_id, &[0.3_f32], Some("m"))
            .unwrap();

        let all = storage.get_all_embeddings().unwrap();
        assert_eq!(all.len(), 3);

        let scoped = storage.get_embeddings_in_buffer(buffer_id).unwrap();
        assert_eq!(scoped.len(), 2);
        assert!(scoped.iter().any(|(chunk_id, _)| *chunk_id == chunk_id1));
        assert!(scoped.iter().any(|(chunk_id, _)| *chunk_id == chunk_id2));
        assert!(
            !scoped
                .iter()
                .any(|(chunk_id, _)| *chunk_id == other_chunk_id)
        );
        assert!(
            storage
                .get_embeddings_in_buffer(other_buffer_id)
                .unwrap()
                .iter()
                .any(|(chunk_id, _)| *chunk_id == other_chunk_id)
        );
    }

    #[test]
    fn test_get_embedding_models() {
        let mut storage = setup();
        let (buffer_id, chunk_id) = setup_buffer_with_chunk(&mut storage);

        // No models initially
        assert!(storage.get_embedding_models(buffer_id).unwrap().is_empty());

        storage
            .store_embedding(chunk_id, &[0.1_f32], Some("model-x"))
            .unwrap();

        let models = storage.get_embedding_models(buffer_id).unwrap();
        assert_eq!(models.len(), 1);
        assert_eq!(models[0], "model-x");
    }

    #[test]
    fn test_get_embedding_model_counts() {
        let mut storage = setup();
        let (buffer_id, chunk_id1) = setup_buffer_with_chunk(&mut storage);

        let chunks2 = vec![Chunk::new(buffer_id, "extra".to_string(), 0..5, 1)];
        storage.add_chunks(buffer_id, &chunks2).unwrap();
        let chunk_id2 = storage
            .get_chunks(buffer_id)
            .unwrap()
            .into_iter()
            .find(|c| c.id != Some(chunk_id1))
            .unwrap()
            .id
            .unwrap();

        storage
            .store_embedding(chunk_id1, &[0.1_f32], Some("model-a"))
            .unwrap();
        storage
            .store_embedding(chunk_id2, &[0.2_f32], Some("model-a"))
            .unwrap();

        let counts = storage.get_embedding_model_counts(buffer_id).unwrap();
        assert_eq!(counts.len(), 1);
        assert_eq!(counts[0], (Some("model-a".to_string()), 2));
    }

    #[test]
    fn test_get_chunks_needing_embedding_no_embeddings() {
        let mut storage = setup();
        let (buffer_id, _chunk_id) = setup_buffer_with_chunk(&mut storage);

        let needing = storage
            .get_chunks_needing_embedding(buffer_id, None)
            .unwrap();
        assert_eq!(needing.len(), 1);
    }

    #[test]
    fn test_get_chunks_needing_embedding_with_model() {
        let mut storage = setup();
        let (buffer_id, chunk_id) = setup_buffer_with_chunk(&mut storage);

        // Add second chunk
        let chunks2 = vec![Chunk::new(buffer_id, "b".to_string(), 0..1, 1)];
        storage.add_chunks(buffer_id, &chunks2).unwrap();
        let chunk_id2 = storage
            .get_chunks(buffer_id)
            .unwrap()
            .into_iter()
            .find(|c| c.id != Some(chunk_id))
            .unwrap()
            .id
            .unwrap();

        // chunk_id has model-a, chunk_id2 has no embedding
        storage
            .store_embedding(chunk_id, &[0.1_f32], Some("model-a"))
            .unwrap();

        // When checking for model-a: chunk_id2 needs one (no embedding)
        let needing = storage
            .get_chunks_needing_embedding(buffer_id, Some("model-a"))
            .unwrap();
        assert!(needing.contains(&chunk_id2));
        assert!(!needing.contains(&chunk_id));

        // When checking for model-b: chunk_id has wrong model, chunk_id2 has none
        let needing_b = storage
            .get_chunks_needing_embedding(buffer_id, Some("model-b"))
            .unwrap();
        assert!(needing_b.contains(&chunk_id));
        assert!(needing_b.contains(&chunk_id2));
    }

    #[test]
    fn test_get_chunks_without_embedding() {
        let mut storage = setup();
        let (buffer_id, chunk_id) = setup_buffer_with_chunk(&mut storage);

        let without = storage.get_chunks_without_embedding(buffer_id).unwrap();
        assert_eq!(without.len(), 1);
        assert!(without.contains(&chunk_id));

        storage.store_embedding(chunk_id, &[0.1_f32], None).unwrap();

        let without_after = storage.get_chunks_without_embedding(buffer_id).unwrap();
        assert!(without_after.is_empty());
    }

    #[test]
    fn test_delete_embeddings_by_model_named() {
        let mut storage = setup();
        let (buffer_id, chunk_id1) = setup_buffer_with_chunk(&mut storage);

        let chunks2 = vec![Chunk::new(buffer_id, "b".to_string(), 0..1, 1)];
        storage.add_chunks(buffer_id, &chunks2).unwrap();
        let chunk_id2 = storage
            .get_chunks(buffer_id)
            .unwrap()
            .into_iter()
            .find(|c| c.id != Some(chunk_id1))
            .unwrap()
            .id
            .unwrap();

        storage
            .store_embedding(chunk_id1, &[0.1_f32], Some("model-a"))
            .unwrap();
        storage
            .store_embedding(chunk_id2, &[0.2_f32], Some("model-b"))
            .unwrap();

        let deleted = storage
            .delete_embeddings_by_model(buffer_id, Some("model-a"))
            .unwrap();
        assert_eq!(deleted, 1);
        assert!(!storage.has_embedding(chunk_id1).unwrap());
        assert!(storage.has_embedding(chunk_id2).unwrap());
    }

    #[test]
    fn test_delete_embeddings_by_model_null() {
        let mut storage = setup();
        let (buffer_id, chunk_id1) = setup_buffer_with_chunk(&mut storage);

        let chunks2 = vec![Chunk::new(buffer_id, "b".to_string(), 0..1, 1)];
        storage.add_chunks(buffer_id, &chunks2).unwrap();
        let chunk_id2 = storage
            .get_chunks(buffer_id)
            .unwrap()
            .into_iter()
            .find(|c| c.id != Some(chunk_id1))
            .unwrap()
            .id
            .unwrap();

        storage
            .store_embedding(chunk_id1, &[0.1_f32], None)
            .unwrap();
        storage
            .store_embedding(chunk_id2, &[0.2_f32], Some("model-b"))
            .unwrap();

        // Delete only the NULL-model embedding
        let deleted = storage.delete_embeddings_by_model(buffer_id, None).unwrap();
        assert_eq!(deleted, 1);
        assert!(!storage.has_embedding(chunk_id1).unwrap());
        assert!(storage.has_embedding(chunk_id2).unwrap());
    }

    #[test]
    fn test_get_embedding_stats() {
        let mut storage = setup();
        let (buffer_id, chunk_id) = setup_buffer_with_chunk(&mut storage);

        // Add second chunk
        let chunks2 = vec![Chunk::new(buffer_id, "extra".to_string(), 0..5, 1)];
        storage.add_chunks(buffer_id, &chunks2).unwrap();

        // Stats before embedding
        let stats = storage.get_embedding_stats(buffer_id).unwrap();
        assert_eq!(stats.total_chunks, 2);
        assert_eq!(stats.embedded_chunks, 0);

        // Embed one chunk
        storage
            .store_embedding(chunk_id, &[0.1_f32], Some("m1"))
            .unwrap();

        let stats = storage.get_embedding_stats(buffer_id).unwrap();
        assert_eq!(stats.total_chunks, 2);
        assert_eq!(stats.embedded_chunks, 1);
        assert_eq!(stats.model_counts.len(), 1);
    }

    #[test]
    fn test_search_fts_finds_match() {
        let mut storage = setup();
        let buffer = Buffer::from_content("The quick brown fox jumps".to_string());
        let buffer_id = storage.add_buffer(&buffer).unwrap();

        let chunks = vec![
            Chunk::new(buffer_id, "The quick brown fox".to_string(), 0..19, 0),
            Chunk::new(buffer_id, "fox jumps".to_string(), 20..29, 1),
        ];
        storage.add_chunks(buffer_id, &chunks).unwrap();

        let results = storage.search_fts("quick", 10).unwrap();
        // At least the first chunk should match
        assert!(!results.is_empty());
        // Scores should be positive
        for (_id, score) in &results {
            assert!(
                *score >= 0.0,
                "BM25 score should be non-negative, got {score}"
            );
        }
    }

    #[test]
    fn test_search_fts_no_match() {
        let mut storage = setup();
        let buffer = Buffer::from_content("The quick brown fox".to_string());
        let buffer_id = storage.add_buffer(&buffer).unwrap();

        let chunks = vec![Chunk::new(
            buffer_id,
            "The quick brown fox".to_string(),
            0..19,
            0,
        )];
        storage.add_chunks(buffer_id, &chunks).unwrap();

        let results = storage.search_fts("zzzyyyxxx", 10).unwrap();
        assert!(results.is_empty());
    }

    #[test]
    fn test_search_fts_respects_limit() {
        let mut storage = setup();
        let buffer = Buffer::from_content("hello world hello world hello".to_string());
        let buffer_id = storage.add_buffer(&buffer).unwrap();

        // Add 3 chunks all containing "hello"
        let chunks = vec![
            Chunk::new(buffer_id, "hello world".to_string(), 0..11, 0),
            Chunk::new(buffer_id, "hello world".to_string(), 12..23, 1),
            Chunk::new(buffer_id, "hello".to_string(), 24..29, 2),
        ];
        storage.add_chunks(buffer_id, &chunks).unwrap();

        let results = storage.search_fts("hello", 2).unwrap();
        assert!(results.len() <= 2);
    }

    #[test]
    fn test_get_chunks_by_ids_empty() {
        let mut storage = SqliteStorage::in_memory().unwrap();
        storage.init().unwrap();

        let result = storage.get_chunks_by_ids(&[]).unwrap();
        assert!(result.is_empty());
    }

    #[test]
    fn test_get_chunks_by_ids_batch() {
        let mut storage = SqliteStorage::in_memory().unwrap();
        storage.init().unwrap();

        let buffer_id = storage
            .add_buffer(&Buffer::from_content("abc def ghi".to_string()))
            .unwrap();
        let chunks = vec![
            Chunk::new(buffer_id, "abc".to_string(), 0..3, 0),
            Chunk::new(buffer_id, "def".to_string(), 4..7, 1),
            Chunk::new(buffer_id, "ghi".to_string(), 8..11, 2),
        ];
        storage.add_chunks(buffer_id, &chunks).unwrap();

        // Fetch by known IDs
        let all = storage.get_chunks(buffer_id).unwrap();
        let ids: Vec<i64> = all.iter().filter_map(|c| c.id).take(2).collect();
        assert_eq!(ids.len(), 2);

        let map = storage.get_chunks_by_ids(&ids).unwrap();
        assert_eq!(map.len(), 2);
        for id in &ids {
            assert!(map.contains_key(id));
        }
    }

    #[test]
    fn test_get_chunks_by_ids_missing_id() {
        let mut storage = SqliteStorage::in_memory().unwrap();
        storage.init().unwrap();

        // Query for an ID that doesn't exist
        let map = storage.get_chunks_by_ids(&[99999]).unwrap();
        assert!(map.is_empty());
    }

    #[test]
    fn test_fts_and_embedding_pipeline() {
        // Integration: index chunks, store embeddings, then verify that FTS
        // hits match chunks that also have embeddings.
        let mut storage = setup();
        let buffer = Buffer::from_content("machine learning and neural networks".to_string());
        let buffer_id = storage.add_buffer(&buffer).unwrap();

        let chunks = vec![
            Chunk::new(
                buffer_id,
                "machine learning algorithms".to_string(),
                0..27,
                0,
            ),
            Chunk::new(
                buffer_id,
                "neural networks architecture".to_string(),
                28..56,
                1,
            ),
        ];
        storage.add_chunks(buffer_id, &chunks).unwrap();

        // Store embeddings for both chunks
        let chunk_ids: Vec<i64> = storage
            .get_chunks(buffer_id)
            .unwrap()
            .into_iter()
            .map(|c| c.id.unwrap())
            .collect();
        assert_eq!(chunk_ids.len(), 2);

        let embeddings: &[&[f32]] = &[&[0.1, 0.2], &[0.2, 0.4]];
        for (&chunk_id, embedding) in chunk_ids.iter().zip(embeddings) {
            storage
                .store_embedding(chunk_id, embedding, Some("test-model"))
                .unwrap();
        }

        // FTS search should find the relevant chunk
        let fts_results = storage.search_fts("machine", 10).unwrap();
        assert!(!fts_results.is_empty());

        // Every FTS result should also have an embedding stored
        for (chunk_id, _score) in &fts_results {
            assert!(
                storage.has_embedding(*chunk_id).unwrap(),
                "FTS result chunk {chunk_id} should have an embedding"
            );
        }

        // Verify embedding count matches what was stored
        assert_eq!(storage.embedding_count().unwrap(), 2);
    }
}