"""Vector storage and similarity search for semantic chunks.

Optimized for high-performance similarity search using:
- HNSW index for O(log N) approximate nearest neighbor search (primary)
- Cached embedding matrix for batch operations (fallback)
- NumPy vectorized cosine similarity (fallback, 100x+ faster than loops)
- Lazy content loading (only fetch for top-k results)
"""

from __future__ import annotations

import json
import logging
import sys
import sqlite3
import threading
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from codexlens.entities import SearchResult, SemanticChunk
from codexlens.errors import StorageError

try:
    import numpy as np
    NUMPY_AVAILABLE = True
except ImportError:
    np = None  # type: ignore[assignment]
    NUMPY_AVAILABLE = False

# Try to import ANN index (optional hnswlib dependency)
try:
    from codexlens.semantic.ann_index import (
        ANNIndex,
        BinaryANNIndex,
        create_ann_index,
        HNSWLIB_AVAILABLE,
    )
except ImportError:
    HNSWLIB_AVAILABLE = False
    ANNIndex = None
    BinaryANNIndex = None
    create_ann_index = None


logger = logging.getLogger(__name__)

# Epsilon used to guard against floating point precision edge cases (e.g., near-zero norms).
EPSILON = 1e-10

# SQLite INTEGER PRIMARY KEY uses signed 64-bit rowids.
SQLITE_INTEGER_MAX = (1 << 63) - 1


def _validate_chunk_id_range(start_id: int, count: int) -> None:
    """Validate that a batch insert can safely generate sequential chunk IDs."""
    if count <= 0:
        return

    last_id = start_id + count - 1
    if last_id > sys.maxsize or last_id > SQLITE_INTEGER_MAX:
        raise ValueError(
            "Chunk ID range overflow: "
            f"start_id={start_id}, count={count} would allocate up to {last_id}, "
            f"exceeding limits (sys.maxsize={sys.maxsize}, sqlite_max={SQLITE_INTEGER_MAX}). "
            "Consider cleaning up the index database or creating a new index database."
        )


def _validate_sql_placeholders(placeholders: str, expected_count: int) -> None:
    """Validate the placeholder string used for a parameterized SQL IN clause."""
    expected = ",".join("?" * expected_count)
    if placeholders != expected:
        raise ValueError(
            "Invalid SQL placeholders for IN clause. "
            f"Expected {expected_count} '?' placeholders."
        )


def _cosine_similarity(a: List[float], b: List[float]) -> float:
    """Compute cosine similarity between two vectors."""
    if not NUMPY_AVAILABLE:
        raise ImportError("numpy required for vector operations")

    a_arr = np.array(a)
    b_arr = np.array(b)

    norm_a = np.linalg.norm(a_arr)
    norm_b = np.linalg.norm(b_arr)

    # Use epsilon tolerance to avoid division by (near-)zero due to floating point precision.
    if norm_a < EPSILON or norm_b < EPSILON:
        return 0.0

    denom = norm_a * norm_b
    if denom < EPSILON:
        return 0.0

    return float(np.dot(a_arr, b_arr) / denom)


class VectorStore:
    """SQLite-based vector storage with HNSW-accelerated similarity search.

    Performance optimizations:
    - HNSW index for O(log N) approximate nearest neighbor search
    - Embedding matrix cached in memory for batch similarity computation (fallback)
    - NumPy vectorized operations instead of Python loops (fallback)
    - Lazy content loading - only fetch full content for top-k results
    - Thread-safe cache invalidation
    - Bulk insert mode for efficient batch operations
    """

    # Default embedding dimension (used when creating new index)
    DEFAULT_DIM = 768

    def __init__(self, db_path: str | Path) -> None:
        if not NUMPY_AVAILABLE:
            raise ImportError(
                "Semantic search dependencies not available. "
                "Install with: pip install codexlens[semantic]"
            )

        self.db_path = Path(db_path)
        self.db_path.parent.mkdir(parents=True, exist_ok=True)

        # Embedding cache for fast similarity search (fallback)
        self._cache_lock = threading.RLock()
        self._embedding_matrix: Optional[np.ndarray] = None
        self._embedding_norms: Optional[np.ndarray] = None
        self._chunk_ids: Optional[List[int]] = None
        self._cache_version: int = 0

        # ANN index for O(log N) search
        self._ann_index: Optional[ANNIndex] = None
        self._ann_dim: Optional[int] = None
        self._ann_write_lock = threading.Lock()  # Protects ANN index modifications

        # Bulk insert mode tracking
        self._bulk_insert_mode: bool = False
        self._bulk_insert_ids: List[int] = []
        self._bulk_insert_embeddings: List[np.ndarray] = []

        self._init_schema()
        self._init_ann_index()

    def _init_schema(self) -> None:
        """Initialize vector storage schema."""
        with sqlite3.connect(self.db_path) as conn:
            # Enable memory mapping for faster reads
            conn.execute("PRAGMA mmap_size = 30000000000")  # 30GB limit
            conn.execute("""
                CREATE TABLE IF NOT EXISTS semantic_chunks (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    file_path TEXT NOT NULL,
                    content TEXT NOT NULL,
                    embedding BLOB NOT NULL,
                    metadata TEXT,
                    category TEXT DEFAULT 'code',
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)
            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_chunks_file
                ON semantic_chunks(file_path)
            """)
            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_chunks_category
                ON semantic_chunks(category)
            """)
            # Model configuration table - tracks which model generated the embeddings
            conn.execute("""
                CREATE TABLE IF NOT EXISTS embeddings_config (
                    id INTEGER PRIMARY KEY CHECK (id = 1),
                    model_profile TEXT NOT NULL,
                    model_name TEXT NOT NULL,
                    embedding_dim INTEGER NOT NULL,
                    backend TEXT NOT NULL DEFAULT 'fastembed',
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)

            # Migration: Add backend column to existing tables
            self._migrate_backend_column(conn)
            # Migration: Add category column
            self._migrate_category_column(conn)

            conn.commit()

    def _migrate_backend_column(self, conn: sqlite3.Connection) -> None:
        """Add backend column to existing embeddings_config table if not present.

        Args:
            conn: Active SQLite connection
        """
        # Check if backend column exists
        cursor = conn.execute("PRAGMA table_info(embeddings_config)")
        columns = [row[1] for row in cursor.fetchall()]

        if 'backend' not in columns:
            logger.info("Migrating embeddings_config table: adding backend column")
            conn.execute("""
                ALTER TABLE embeddings_config
                ADD COLUMN backend TEXT NOT NULL DEFAULT 'fastembed'
            """)

    def _migrate_category_column(self, conn: sqlite3.Connection) -> None:
        """Add category column to existing semantic_chunks table if not present.

        Args:
            conn: Active SQLite connection
        """
        # Check if category column exists
        cursor = conn.execute("PRAGMA table_info(semantic_chunks)")
        columns = [row[1] for row in cursor.fetchall()]

        if 'category' not in columns:
            logger.info("Migrating semantic_chunks table: adding category column")
            conn.execute("""
                ALTER TABLE semantic_chunks
                ADD COLUMN category TEXT DEFAULT 'code'
            """)
            # Create index for fast category filtering
            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_chunks_category
                ON semantic_chunks(category)
            """)

    def _init_ann_index(self) -> None:
        """Initialize ANN index (lazy loading from existing data)."""
        if not HNSWLIB_AVAILABLE:
            logger.debug("hnswlib not available, using brute-force search")
            return

        # Try to detect embedding dimension from existing data
        dim = self._detect_embedding_dim()
        if dim is None:
            # No data yet, will initialize on first add
            logger.debug("No embeddings found, ANN index will be created on first add")
            return

        self._ann_dim = dim

        try:
            self._ann_index = ANNIndex(self.db_path, dim)
            if self._ann_index.load():
                logger.debug(
                    "Loaded ANN index with %d vectors", self._ann_index.count()
                )
            else:
                # Index file doesn't exist, try to build from SQLite data
                logger.debug("ANN index file not found, rebuilding from SQLite")
                self._rebuild_ann_index_internal()
        except Exception as e:
            logger.warning("Failed to initialize ANN index: %s", e)
            self._ann_index = None

    def _detect_embedding_dim(self) -> Optional[int]:
        """Detect embedding dimension from existing data."""
        with sqlite3.connect(self.db_path) as conn:
            row = conn.execute(
                "SELECT embedding FROM semantic_chunks LIMIT 1"
            ).fetchone()
            if row and row[0]:
                # Embedding is stored as float32 blob
                blob = row[0]
                return len(blob) // np.dtype(np.float32).itemsize
        return None

    @property
    def dimension(self) -> Optional[int]:
        """Return the dimension of embeddings in the store.

        Returns:
            Embedding dimension if available, None if store is empty.
        """
        if self._ann_dim is not None:
            return self._ann_dim
        self._ann_dim = self._detect_embedding_dim()
        return self._ann_dim

    def _rebuild_ann_index_internal(self) -> int:
        """Internal method to rebuild ANN index from SQLite data."""
        if self._ann_index is None:
            return 0

        with sqlite3.connect(self.db_path) as conn:
            conn.execute("PRAGMA mmap_size = 30000000000")
            rows = conn.execute(
                "SELECT id, embedding FROM semantic_chunks"
            ).fetchall()

        if not rows:
            return 0

        # Extract IDs and embeddings
        ids = [r[0] for r in rows]
        embeddings = np.vstack([
            np.frombuffer(r[1], dtype=np.float32) for r in rows
        ])

        # Add to ANN index
        self._ann_index.add_vectors(ids, embeddings)
        self._ann_index.save()

        logger.info("Rebuilt ANN index with %d vectors", len(ids))
        return len(ids)

    def rebuild_ann_index(self) -> int:
        """Rebuild HNSW index from all chunks in SQLite.

        Use this method to:
        - Migrate existing data to use ANN search
        - Repair corrupted index
        - Reclaim space after many deletions

        Returns:
            Number of vectors indexed.
        """
        if not HNSWLIB_AVAILABLE:
            logger.warning("hnswlib not available, cannot rebuild ANN index")
            return 0

        # Detect dimension
        dim = self._detect_embedding_dim()
        if dim is None:
            logger.warning("No embeddings found, cannot rebuild ANN index")
            return 0

        self._ann_dim = dim

        # Create new index
        try:
            self._ann_index = ANNIndex(self.db_path, dim)
            return self._rebuild_ann_index_internal()
        except Exception as e:
            logger.error("Failed to rebuild ANN index: %s", e)
            self._ann_index = None
            return 0

    def _invalidate_cache(self) -> None:
        """Invalidate the embedding cache (thread-safe)."""
        with self._cache_lock:
            self._embedding_matrix = None
            self._embedding_norms = None
            self._chunk_ids = None
            self._cache_version += 1

    def _refresh_cache(self) -> bool:
        """Load embeddings into numpy matrix for fast similarity search.

        Returns:
            True if cache was refreshed successfully, False if no data.
        """
        with self._cache_lock:
            with sqlite3.connect(self.db_path) as conn:
                conn.execute("PRAGMA mmap_size = 30000000000")
                rows = conn.execute(
                    "SELECT id, embedding FROM semantic_chunks"
                ).fetchall()

            if not rows:
                self._embedding_matrix = None
                self._embedding_norms = None
                self._chunk_ids = None
                return False

            # Extract IDs and embeddings
            self._chunk_ids = [r[0] for r in rows]

            # Bulk convert binary blobs to numpy matrix
            embeddings = [
                np.frombuffer(r[1], dtype=np.float32) for r in rows
            ]
            self._embedding_matrix = np.vstack(embeddings)

            # Pre-compute norms for faster similarity calculation
            self._embedding_norms = np.linalg.norm(
                self._embedding_matrix, axis=1, keepdims=True
            )
            # Avoid division by zero
            self._embedding_norms = np.where(
                self._embedding_norms == 0, EPSILON, self._embedding_norms
            )

            return True

    def _ensure_ann_index(self, dim: int) -> bool:
        """Ensure ANN index is initialized with correct dimension.

        This method is thread-safe and uses double-checked locking.

        Args:
            dim: Embedding dimension

        Returns:
            True if ANN index is ready, False otherwise
        """
        if not HNSWLIB_AVAILABLE:
            return False

        # Fast path: index already initialized (no lock needed)
        if self._ann_index is not None:
            return True

        # Slow path: acquire lock for initialization
        with self._ann_write_lock:
            # Double-check after acquiring lock
            if self._ann_index is not None:
                return True

            try:
                self._ann_dim = dim
                self._ann_index = ANNIndex(self.db_path, dim)
                self._ann_index.load()  # Try to load existing
                return True
            except Exception as e:
                logger.warning("Failed to initialize ANN index: %s", e)
                self._ann_index = None
                return False

    def add_chunk(
        self, chunk: SemanticChunk, file_path: str, category: str = "code"
    ) -> int:
        """Add a single chunk with its embedding.

        Args:
            chunk: SemanticChunk with embedding
            file_path: Path to the source file
            category: File category ('code' or 'doc'), default 'code'

        Returns:
            The inserted chunk ID.
        """
        if chunk.embedding is None:
            raise ValueError("Chunk must have embedding before adding to store")

        embedding_arr = np.array(chunk.embedding, dtype=np.float32)
        embedding_blob = embedding_arr.tobytes()
        metadata_json = json.dumps(chunk.metadata) if chunk.metadata else None

        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.execute(
                """
                INSERT INTO semantic_chunks (file_path, content, embedding, metadata, category)
                VALUES (?, ?, ?, ?, ?)
                """,
                (file_path, chunk.content, embedding_blob, metadata_json, category)
            )
            conn.commit()
            chunk_id = cursor.lastrowid or 0

        # Add to ANN index
        if self._ensure_ann_index(len(chunk.embedding)):
            with self._ann_write_lock:
                try:
                    self._ann_index.add_vectors([chunk_id], embedding_arr.reshape(1, -1))
                    self._ann_index.save()
                except Exception as e:
                    logger.warning("Failed to add to ANN index: %s", e)

        # Invalidate cache after modification
        self._invalidate_cache()
        return chunk_id

    def add_chunks(
        self, chunks: List[SemanticChunk], file_path: str, category: str = "code"
    ) -> List[int]:
        """Add multiple chunks with embeddings (batch insert).

        Args:
            chunks: List of SemanticChunk objects with embeddings
            file_path: Path to the source file
            category: File category ('code' or 'doc'), default 'code'

        Returns:
            List of inserted chunk IDs.
        """
        if not chunks:
            return []

        # Prepare batch data
        batch_data = []
        embeddings_list = []
        for chunk in chunks:
            if chunk.embedding is None:
                raise ValueError("All chunks must have embeddings")
            embedding_arr = np.array(chunk.embedding, dtype=np.float32)
            embedding_blob = embedding_arr.tobytes()
            metadata_json = json.dumps(chunk.metadata) if chunk.metadata else None
            batch_data.append((file_path, chunk.content, embedding_blob, metadata_json, category))
            embeddings_list.append(embedding_arr)

        # Batch insert to SQLite
        with sqlite3.connect(self.db_path) as conn:
            # Get starting ID before insert
            row = conn.execute("SELECT MAX(id) FROM semantic_chunks").fetchone()
            start_id = (row[0] or 0) + 1

            conn.executemany(
                """
                INSERT INTO semantic_chunks (file_path, content, embedding, metadata, category)
                VALUES (?, ?, ?, ?, ?)
                """,
                batch_data
            )
            conn.commit()
            # Calculate inserted IDs based on starting ID
            ids = list(range(start_id, start_id + len(chunks)))

        # Add to ANN index
        if embeddings_list and self._ensure_ann_index(len(embeddings_list[0])):
            with self._ann_write_lock:
                try:
                    embeddings_matrix = np.vstack(embeddings_list)
                    self._ann_index.add_vectors(ids, embeddings_matrix)
                    self._ann_index.save()
                except Exception as e:
                    logger.warning("Failed to add batch to ANN index: %s", e)

        # Invalidate cache after modification
        self._invalidate_cache()
        return ids

    def add_chunks_batch(
        self,
        chunks_with_paths: List[Tuple[SemanticChunk, str]],
        update_ann: bool = True,
        auto_save_ann: bool = True,
        categories: Optional[List[str]] = None,
    ) -> List[int]:
        """Batch insert chunks from multiple files in a single transaction.

        This method is optimized for bulk operations during index generation.

        Args:
            chunks_with_paths: List of (chunk, file_path) tuples
            update_ann: If True, update ANN index with new vectors (default: True)
            auto_save_ann: If True, save ANN index after update (default: True).
                          Set to False for bulk inserts to reduce I/O overhead.
            categories: Optional list of categories per chunk. If None, defaults to 'code'.
                       If provided, must match length of chunks_with_paths.

        Returns:
            List of inserted chunk IDs
        """
        if not chunks_with_paths:
            return []

        batch_size = len(chunks_with_paths)

        # Validate categories if provided
        if categories is not None and len(categories) != batch_size:
            raise ValueError(
                f"categories length ({len(categories)}) must match "
                f"chunks_with_paths length ({batch_size})"
            )

        # Prepare batch data
        batch_data = []
        embeddings_list = []
        for i, (chunk, file_path) in enumerate(chunks_with_paths):
            if chunk.embedding is None:
                raise ValueError("All chunks must have embeddings")
            # Optimize: avoid repeated np.array() if already numpy
            if isinstance(chunk.embedding, np.ndarray):
                embedding_arr = chunk.embedding.astype(np.float32)
            else:
                embedding_arr = np.array(chunk.embedding, dtype=np.float32)
            embedding_blob = embedding_arr.tobytes()
            metadata_json = json.dumps(chunk.metadata) if chunk.metadata else None
            category = categories[i] if categories else "code"
            batch_data.append((file_path, chunk.content, embedding_blob, metadata_json, category))
            embeddings_list.append(embedding_arr)

        # Batch insert to SQLite in single transaction
        with sqlite3.connect(self.db_path) as conn:
            # Get starting ID before insert
            row = conn.execute("SELECT MAX(id) FROM semantic_chunks").fetchone()
            start_id = (row[0] or 0) + 1

            _validate_chunk_id_range(start_id, batch_size)

            conn.executemany(
                """
                INSERT INTO semantic_chunks (file_path, content, embedding, metadata, category)
                VALUES (?, ?, ?, ?, ?)
                """,
                batch_data
            )
            conn.commit()
            # Calculate inserted IDs based on starting ID
            ids = list(range(start_id, start_id + batch_size))

        # Handle ANN index updates
        if embeddings_list and update_ann and self._ensure_ann_index(len(embeddings_list[0])):
            with self._ann_write_lock:
                # In bulk insert mode, accumulate for later batch update
                if self._bulk_insert_mode:
                    self._bulk_insert_ids.extend(ids)
                    self._bulk_insert_embeddings.extend(embeddings_list)
                else:
                    # Normal mode: update immediately
                    try:
                        embeddings_matrix = np.vstack(embeddings_list)
                        self._ann_index.add_vectors(ids, embeddings_matrix)
                        if auto_save_ann:
                            self._ann_index.save()
                    except Exception as e:
                        logger.warning("Failed to add batch to ANN index: %s", e)

        # Invalidate cache after modification
        self._invalidate_cache()
        return ids

    def add_chunks_batch_numpy(
        self,
        chunks_with_paths: List[Tuple[SemanticChunk, str]],
        embeddings_matrix: np.ndarray,
        update_ann: bool = True,
        auto_save_ann: bool = True,
        categories: Optional[List[str]] = None,
    ) -> List[int]:
        """Batch insert chunks with pre-computed numpy embeddings matrix.

        This method accepts embeddings as a numpy matrix to avoid list->array conversions.
        Useful when embeddings are already in numpy format from batch encoding.

        Args:
            chunks_with_paths: List of (chunk, file_path) tuples (embeddings can be None)
            embeddings_matrix: Pre-computed embeddings as (N, D) numpy array
            update_ann: If True, update ANN index with new vectors (default: True)
            auto_save_ann: If True, save ANN index after update (default: True)
            categories: Optional list of categories per chunk. If None, defaults to 'code'.

        Returns:
            List of inserted chunk IDs
        """
        if not chunks_with_paths:
            return []

        batch_size = len(chunks_with_paths)

        if len(chunks_with_paths) != embeddings_matrix.shape[0]:
            raise ValueError(
                f"Mismatch: {len(chunks_with_paths)} chunks but "
                f"{embeddings_matrix.shape[0]} embeddings"
            )

        # Validate categories if provided
        if categories is not None and len(categories) != batch_size:
            raise ValueError(
                f"categories length ({len(categories)}) must match "
                f"chunks_with_paths length ({batch_size})"
            )

        # Ensure float32 format
        embeddings_matrix = embeddings_matrix.astype(np.float32)

        # Prepare batch data
        batch_data = []
        for i, (chunk, file_path) in enumerate(chunks_with_paths):
            embedding_arr = embeddings_matrix[i]
            embedding_blob = embedding_arr.tobytes()
            metadata_json = json.dumps(chunk.metadata) if chunk.metadata else None
            category = categories[i] if categories else "code"
            batch_data.append((file_path, chunk.content, embedding_blob, metadata_json, category))

        # Batch insert to SQLite in single transaction
        with sqlite3.connect(self.db_path) as conn:
            # Get starting ID before insert
            row = conn.execute("SELECT MAX(id) FROM semantic_chunks").fetchone()
            start_id = (row[0] or 0) + 1

            _validate_chunk_id_range(start_id, batch_size)

            conn.executemany(
                """
                INSERT INTO semantic_chunks (file_path, content, embedding, metadata, category)
                VALUES (?, ?, ?, ?, ?)
                """,
                batch_data
            )
            conn.commit()
            # Calculate inserted IDs based on starting ID
            ids = list(range(start_id, start_id + batch_size))

        # Handle ANN index updates
        if update_ann and self._ensure_ann_index(embeddings_matrix.shape[1]):
            with self._ann_write_lock:
                # In bulk insert mode, accumulate for later batch update
                if self._bulk_insert_mode:
                    self._bulk_insert_ids.extend(ids)
                    # Split matrix into individual arrays for accumulation
                    self._bulk_insert_embeddings.extend([embeddings_matrix[i] for i in range(len(ids))])
                else:
                    # Normal mode: update immediately
                    try:
                        self._ann_index.add_vectors(ids, embeddings_matrix)
                        if auto_save_ann:
                            self._ann_index.save()
                    except Exception as e:
                        logger.warning("Failed to add batch to ANN index: %s", e)

        # Invalidate cache after modification
        self._invalidate_cache()
        return ids

    def begin_bulk_insert(self) -> None:
        """Begin bulk insert mode - disable ANN auto-update for better performance.

        Usage:
            store.begin_bulk_insert()
            try:
                for batch in batches:
                    store.add_chunks_batch(batch, auto_save_ann=False)
            finally:
                store.end_bulk_insert()

        Or use context manager:
            with store.bulk_insert():
                for batch in batches:
                    store.add_chunks_batch(batch)
        """
        with self._ann_write_lock:
            self._bulk_insert_mode = True
            self._bulk_insert_ids.clear()
            self._bulk_insert_embeddings.clear()
        logger.debug("Entered bulk insert mode")

    def end_bulk_insert(self) -> None:
        """End bulk insert mode and rebuild ANN index from accumulated data.

        This method should be called after all bulk inserts are complete to
        update the ANN index in a single batch operation.
        """
        with self._ann_write_lock:
            if not self._bulk_insert_mode:
                logger.warning("end_bulk_insert called but not in bulk insert mode")
                return

            self._bulk_insert_mode = False
            bulk_ids = list(self._bulk_insert_ids)
            bulk_embeddings = list(self._bulk_insert_embeddings)
            self._bulk_insert_ids.clear()
            self._bulk_insert_embeddings.clear()

        # Update ANN index with accumulated data.
        if bulk_ids and bulk_embeddings:
            if self._ensure_ann_index(len(bulk_embeddings[0])):
                with self._ann_write_lock:
                    try:
                        embeddings_matrix = np.vstack(bulk_embeddings)
                        self._ann_index.add_vectors(bulk_ids, embeddings_matrix)
                        self._ann_index.save()
                        logger.info(
                            "Bulk insert complete: added %d vectors to ANN index",
                            len(bulk_ids),
                        )
                    except Exception as e:
                        logger.error("Failed to update ANN index after bulk insert: %s", e)

        logger.debug("Exited bulk insert mode")

    class BulkInsertContext:
        """Context manager for bulk insert operations."""

        def __init__(self, store: "VectorStore") -> None:
            self.store = store

        def __enter__(self) -> "VectorStore":
            self.store.begin_bulk_insert()
            return self.store

        def __exit__(self, exc_type, exc_val, exc_tb) -> None:
            self.store.end_bulk_insert()

    def bulk_insert(self) -> "VectorStore.BulkInsertContext":
        """Return a context manager for bulk insert operations.

        Usage:
            with store.bulk_insert():
                for batch in batches:
                    store.add_chunks_batch(batch)
        """
        return self.BulkInsertContext(self)

    def delete_file_chunks(self, file_path: str) -> int:
        """Delete all chunks for a file.

        Returns:
            Number of deleted chunks.
        """
        # Get chunk IDs before deletion (for ANN index)
        chunk_ids_to_delete = []
        if self._ann_index is not None:
            with sqlite3.connect(self.db_path) as conn:
                rows = conn.execute(
                    "SELECT id FROM semantic_chunks WHERE file_path = ?",
                    (file_path,)
                ).fetchall()
                chunk_ids_to_delete = [r[0] for r in rows]

        # Delete from SQLite
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.execute(
                "DELETE FROM semantic_chunks WHERE file_path = ?",
                (file_path,)
            )
            conn.commit()
            deleted = cursor.rowcount

        # Remove from ANN index
        if deleted > 0 and self._ann_index is not None and chunk_ids_to_delete:
            with self._ann_write_lock:
                try:
                    self._ann_index.remove_vectors(chunk_ids_to_delete)
                    self._ann_index.save()
                except Exception as e:
                    logger.warning("Failed to remove from ANN index: %s", e)

        if deleted > 0:
            self._invalidate_cache()
        return deleted

    def search_similar(
        self,
        query_embedding: List[float],
        top_k: int = 10,
        min_score: float = 0.0,
        return_full_content: bool = True,
        category: Optional[str] = None,
    ) -> List[SearchResult]:
        """Find chunks most similar to query embedding.

        Uses HNSW index for O(log N) search when available, falls back to
        brute-force NumPy search otherwise.

        Args:
            query_embedding: Query vector.
            top_k: Maximum results to return.
            min_score: Minimum cosine similarity score in [0.0, 1.0].
            return_full_content: If True, return full code block content.
            category: Optional category filter ('code' or 'doc'). If None, returns all.

        Returns:
            List of SearchResult ordered by similarity (highest first).
        """
        query_vec = np.array(query_embedding, dtype=np.float32)

        if not 0.0 <= min_score <= 1.0:
            raise ValueError(
                f"Invalid min_score: {min_score}. Must be within [0.0, 1.0] for cosine similarity."
            )

        # Try HNSW search first (O(log N))
        if (
            HNSWLIB_AVAILABLE
            and self._ann_index is not None
            and self._ann_index.is_loaded
            and self._ann_index.count() > 0
        ):
            try:
                return self._search_with_ann(
                    query_vec, top_k, min_score, return_full_content, category
                )
            except Exception as e:
                logger.warning("ANN search failed, falling back to brute-force: %s", e)

        # Fallback to brute-force search (O(N))
        return self._search_brute_force(
            query_vec, top_k, min_score, return_full_content, category
        )

    def _search_with_ann(
        self,
        query_vec: np.ndarray,
        top_k: int,
        min_score: float,
        return_full_content: bool,
        category: Optional[str] = None,
    ) -> List[SearchResult]:
        """Search using HNSW index (O(log N)).

        Args:
            query_vec: Query vector as numpy array
            top_k: Maximum results to return
            min_score: Minimum cosine similarity score in [0.0, 1.0]
            return_full_content: If True, return full code block content
            category: Optional category filter ('code' or 'doc')

        Returns:
            List of SearchResult ordered by similarity (highest first)
        """
        # Limit top_k to available vectors to prevent hnswlib error
        ann_count = self._ann_index.count()
        # When category filtering, fetch more candidates to compensate for filtering
        fetch_k = top_k * 3 if category else top_k
        effective_top_k = min(fetch_k, ann_count) if ann_count > 0 else 0

        if effective_top_k == 0:
            return []

        # HNSW search returns (ids, distances)
        # For cosine space: distance = 1 - similarity
        ids, distances = self._ann_index.search(query_vec, effective_top_k)

        if ids is None or distances is None:
            logger.debug(
                "ANN search returned null results (ids=%s, distances=%s)",
                ids,
                distances,
            )
            return []

        if len(ids) == 0 or len(distances) == 0:
            logger.debug(
                "ANN search returned empty results (ids=%s, distances=%s)",
                ids,
                distances,
            )
            return []

        if len(ids) != len(distances):
            logger.warning(
                "ANN search returned mismatched result lengths (%d ids, %d distances)",
                len(ids),
                len(distances),
            )
            return []

        # Convert distances to similarity scores
        scores = [1.0 - d for d in distances]

        # Filter by min_score
        filtered = [
            (chunk_id, score)
            for chunk_id, score in zip(ids, scores)
            if score >= min_score
        ]

        if not filtered:
            return []

        top_ids = [f[0] for f in filtered]
        top_scores = [f[1] for f in filtered]

        # Fetch content from SQLite with category filtering
        results = self._fetch_results_by_ids(
            top_ids, top_scores, return_full_content, category
        )
        # Apply final limit after category filtering
        return results[:top_k]

    def _search_brute_force(
        self,
        query_vec: np.ndarray,
        top_k: int,
        min_score: float,
        return_full_content: bool,
        category: Optional[str] = None,
    ) -> List[SearchResult]:
        """Brute-force search using NumPy (O(N) fallback).

        Args:
            query_vec: Query vector as numpy array
            top_k: Maximum results to return
            min_score: Minimum cosine similarity score in [0.0, 1.0]
            return_full_content: If True, return full code block content
            category: Optional category filter ('code' or 'doc')

        Returns:
            List of SearchResult ordered by similarity (highest first)
        """
        logger.warning(
            "Using brute-force vector search (hnswlib not available). "
            "This may cause high memory usage for large indexes. "
            "Install hnswlib for better performance: pip install hnswlib"
        )

        with self._cache_lock:
            # Refresh cache if needed
            if self._embedding_matrix is None:
                if not self._refresh_cache():
                    return []  # No data

            # Vectorized cosine similarity
            query_vec = query_vec.reshape(1, -1)
            query_norm = np.linalg.norm(query_vec)
            if query_norm == 0:
                return []

            # Compute all similarities at once: (N,) scores
            # similarity = (A @ B.T) / (||A|| * ||B||)
            dot_products = np.dot(self._embedding_matrix, query_vec.T).flatten()
            scores = dot_products / (self._embedding_norms.flatten() * query_norm)

            # Filter by min_score and get top-k indices
            valid_mask = scores >= min_score
            valid_indices = np.where(valid_mask)[0]

            if len(valid_indices) == 0:
                return []

            # When category filtering, fetch more candidates to compensate for filtering
            fetch_k = top_k * 3 if category else top_k

            # Sort by score descending and take top candidates
            valid_scores = scores[valid_indices]
            sorted_order = np.argsort(valid_scores)[::-1][:fetch_k]
            top_indices = valid_indices[sorted_order]
            top_scores = valid_scores[sorted_order]

            # Get chunk IDs for top results
            top_ids = [self._chunk_ids[i] for i in top_indices]

        # Fetch content only for top-k results (lazy loading) with category filtering
        results = self._fetch_results_by_ids(
            top_ids, top_scores.tolist(), return_full_content, category
        )
        # Apply final limit after category filtering
        return results[:top_k]

    def _fetch_results_by_ids(
        self,
        chunk_ids: List[int],
        scores: List[float],
        return_full_content: bool,
        category: Optional[str] = None,
    ) -> List[SearchResult]:
        """Fetch full result data for specific chunk IDs.

        Args:
            chunk_ids: List of chunk IDs to fetch.
            scores: Corresponding similarity scores.
            return_full_content: Whether to include full content.
            category: Optional category filter ('code' or 'doc').

        Returns:
            List of SearchResult objects.
        """
        if not chunk_ids:
            return []

        # Build parameterized query for IN clause
        placeholders = ",".join("?" * len(chunk_ids))
        _validate_sql_placeholders(placeholders, len(chunk_ids))

        # SQL injection prevention:
        # - Only a validated placeholders string (commas + '?') is interpolated into the query.
        # - User-provided values are passed separately via sqlite3 parameters.
        # - Category filter is added as a separate parameter
        if category:
            query = """
                SELECT id, file_path, content, metadata
                FROM semantic_chunks
                WHERE id IN ({placeholders}) AND category = ?
            """.format(placeholders=placeholders)
            params = list(chunk_ids) + [category]
        else:
            query = """
                SELECT id, file_path, content, metadata
                FROM semantic_chunks
                WHERE id IN ({placeholders})
            """.format(placeholders=placeholders)
            params = chunk_ids

        with sqlite3.connect(self.db_path) as conn:
            conn.execute("PRAGMA mmap_size = 30000000000")
            rows = conn.execute(query, params).fetchall()

        # Build ID -> row mapping
        id_to_row = {r[0]: r for r in rows}

        results = []
        for chunk_id, score in zip(chunk_ids, scores):
            row = id_to_row.get(chunk_id)
            if not row:
                continue

            _, file_path, content, metadata_json = row
            metadata = json.loads(metadata_json) if metadata_json else {}

            # Build excerpt (short preview)
            excerpt = content[:200] + "..." if len(content) > 200 else content

            # Extract symbol information from metadata
            symbol_name = metadata.get("symbol_name")
            symbol_kind = metadata.get("symbol_kind")
            start_line = metadata.get("start_line")
            end_line = metadata.get("end_line")

            # Build Symbol object if we have symbol info
            symbol = None
            if symbol_name and symbol_kind and start_line and end_line:
                try:
                    from codexlens.entities import Symbol
                    symbol = Symbol(
                        name=symbol_name,
                        kind=symbol_kind,
                        range=(start_line, end_line)
                    )
                except Exception:
                    pass

            results.append(SearchResult(
                path=file_path,
                score=score,
                excerpt=excerpt,
                content=content if return_full_content else None,
                symbol=symbol,
                metadata=metadata,
                start_line=start_line,
                end_line=end_line,
                symbol_name=symbol_name,
                symbol_kind=symbol_kind,
            ))

        return results

    def count_chunks(self) -> int:
        """Count total chunks in store."""
        with sqlite3.connect(self.db_path) as conn:
            row = conn.execute("SELECT COUNT(*) FROM semantic_chunks").fetchone()
            return row[0] if row else 0

    def get_all_chunks(self) -> List[SemanticChunk]:
        """Get all chunks from the store.

        Returns:
            List of SemanticChunk objects with id and content.
        """
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            rows = conn.execute(
                "SELECT id, file_path, content, metadata FROM semantic_chunks"
            ).fetchall()

            chunks = []
            for row in rows:
                chunks.append(SemanticChunk(
                    id=row["id"],
                    content=row["content"],
                    file_path=row["file_path"],
                    metadata=json.loads(row["metadata"]) if row["metadata"] else None,
                ))
            return chunks

    def clear_cache(self) -> None:
        """Manually clear the embedding cache."""
        self._invalidate_cache()

    @property
    def ann_available(self) -> bool:
        """Check if ANN index is available and ready."""
        return (
            HNSWLIB_AVAILABLE
            and self._ann_index is not None
            and self._ann_index.is_loaded
        )

    @property
    def ann_count(self) -> int:
        """Get number of vectors in ANN index."""
        if self._ann_index is not None:
            return self._ann_index.count()
        return 0

    def get_model_config(self) -> Optional[Dict[str, Any]]:
        """Get the model configuration used for embeddings in this store.

        Returns:
            Dictionary with model_profile, model_name, embedding_dim, backend, or None if not set.
        """
        with sqlite3.connect(self.db_path) as conn:
            row = conn.execute(
                "SELECT model_profile, model_name, embedding_dim, backend, created_at, updated_at "
                "FROM embeddings_config WHERE id = 1"
            ).fetchone()
            if row:
                return {
                    "model_profile": row[0],
                    "model_name": row[1],
                    "embedding_dim": row[2],
                    "backend": row[3],
                    "created_at": row[4],
                    "updated_at": row[5],
                }
        return None

    def set_model_config(
        self, model_profile: str, model_name: str, embedding_dim: int, backend: str = 'fastembed'
    ) -> None:
        """Set the model configuration for embeddings in this store.

        This should be called when generating new embeddings. If a different
        model was previously used, this will update the configuration.

        Args:
            model_profile: Model profile name (fast, code, minilm, etc.)
            model_name: Full model name (e.g., jinaai/jina-embeddings-v2-base-code)
            embedding_dim: Embedding dimension (e.g., 768)
            backend: Backend used for embeddings (fastembed or litellm, default: fastembed)
        """
        with sqlite3.connect(self.db_path) as conn:
            conn.execute(
                """
                INSERT INTO embeddings_config (id, model_profile, model_name, embedding_dim, backend)
                VALUES (1, ?, ?, ?, ?)
                ON CONFLICT(id) DO UPDATE SET
                    model_profile = excluded.model_profile,
                    model_name = excluded.model_name,
                    embedding_dim = excluded.embedding_dim,
                    backend = excluded.backend,
                    updated_at = CURRENT_TIMESTAMP
                """,
                (model_profile, model_name, embedding_dim, backend)
            )
            conn.commit()

    def check_model_compatibility(
        self, model_profile: str, model_name: str, embedding_dim: int
    ) -> Tuple[bool, Optional[str]]:
        """Check if the given model is compatible with existing embeddings.

        Args:
            model_profile: Model profile to check
            model_name: Model name to check
            embedding_dim: Embedding dimension to check

        Returns:
            Tuple of (is_compatible, warning_message).
            is_compatible is True if no existing config or configs match.
            warning_message is a user-friendly message if incompatible.
        """
        existing = self.get_model_config()
        if existing is None:
            return True, None

        # Check dimension first (most critical)
        if existing["embedding_dim"] != embedding_dim:
            return False, (
                f"Dimension mismatch: existing embeddings use {existing['embedding_dim']}d "
                f"({existing['model_profile']}), but requested model uses {embedding_dim}d "
                f"({model_profile}). Use --force to regenerate all embeddings."
            )

        # Check model (different models with same dimension may have different semantic spaces)
        if existing["model_profile"] != model_profile:
            return False, (
                f"Model mismatch: existing embeddings use '{existing['model_profile']}' "
                f"({existing['model_name']}), but requested '{model_profile}' "
                f"({model_name}). Use --force to regenerate all embeddings."
            )

        return True, None

    def close(self) -> None:
        """Close the vector store and release resources.

        This ensures SQLite connections are closed and ANN index is cleared,
        allowing temporary files to be deleted on Windows.
        """
        with self._cache_lock:
            self._embedding_matrix = None
            self._embedding_norms = None
            self._chunk_ids = None

        with self._ann_write_lock:
            self._ann_index = None

    def __enter__(self) -> "VectorStore":
        """Context manager entry."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        """Context manager exit - close resources."""
        self.close()
