"""
Standard Retrieval (RAG) tool with embedding-based search.

This module provides a retrieval tool that uses embeddings to search
through a knowledge base and return relevant documents/chunks.

Features:
- Multiple embedding providers (Sentence Transformers, TF-IDF fallback)
- Document loaders: txt, md, pdf, csv, tsv, json, jsonl
- Chunking strategies: fixed-size, sentence-based, paragraph-based, recursive
- Hybrid search: vector similarity + BM25 keyword matching
- Persistent index storage
"""

import csv
import math
import os
import json
import pickle
import logging
import hashlib
import re
from collections import Counter
from pathlib import Path
from typing import Dict, Any, List, Optional
from dataclasses import dataclass, field
import numpy as np

from ..base_tool import (
    BaseTool,
    ToolMetadata,
    ToolCategory,
    ParameterSpec,
    ParameterType,
)


logger = logging.getLogger(__name__)


@dataclass
class Document:
    """A document in the knowledge base."""
    id: str
    content: str
    metadata: Dict[str, Any] = field(default_factory=dict)
    embedding: Optional[np.ndarray] = None


@dataclass
class RetrievalResult:
    """A retrieval result with score."""
    document: Document
    score: float
    rank: int


# ---------------------------------------------------------------------------
# Embedding Providers
# ---------------------------------------------------------------------------

class EmbeddingProvider:
    """Base class for embedding providers."""

    def embed(self, texts: List[str]) -> np.ndarray:
        """Embed a list of texts."""
        raise NotImplementedError

    def embed_query(self, query: str) -> np.ndarray:
        """Embed a query text."""
        return self.embed([query])[0]


class SentenceTransformerEmbedding(EmbeddingProvider):
    """Sentence Transformers embedding provider."""

    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
        self.model_name = model_name
        self._model = None

    def _load_model(self):
        if self._model is None:
            try:
                from sentence_transformers import SentenceTransformer
                self._model = SentenceTransformer(self.model_name)
                logger.info(f"Loaded embedding model: {self.model_name}")
            except ImportError:
                raise RuntimeError(
                    "sentence-transformers not installed. "
                    "Install with: pip install sentence-transformers"
                )
        return self._model

    def embed(self, texts: List[str]) -> np.ndarray:
        model = self._load_model()
        embeddings = model.encode(texts, convert_to_numpy=True)
        return embeddings


class SimpleEmbedding(EmbeddingProvider):
    """Simple TF-IDF based embedding for environments without sentence-transformers."""

    def __init__(self):
        self._vectorizer = None
        self._fitted = False

    def _get_vectorizer(self):
        if self._vectorizer is None:
            try:
                from sklearn.feature_extraction.text import TfidfVectorizer
                self._vectorizer = TfidfVectorizer(
                    max_features=512,
                    stop_words='english',
                    ngram_range=(1, 2)
                )
            except ImportError:
                raise RuntimeError(
                    "scikit-learn not installed. "
                    "Install with: pip install scikit-learn"
                )
        return self._vectorizer

    def fit(self, texts: List[str]):
        vectorizer = self._get_vectorizer()
        vectorizer.fit(texts)
        self._fitted = True

    def embed(self, texts: List[str]) -> np.ndarray:
        vectorizer = self._get_vectorizer()
        if not self._fitted:
            embeddings = vectorizer.fit_transform(texts).toarray()
            self._fitted = True
        else:
            embeddings = vectorizer.transform(texts).toarray()
        return embeddings


# ---------------------------------------------------------------------------
# Chunking Strategies
# ---------------------------------------------------------------------------

class ChunkingStrategy:
    """Base chunking strategy."""

    def chunk(self, text: str, doc_id: str) -> List[Document]:
        raise NotImplementedError


class FixedSizeChunker(ChunkingStrategy):
    """Fixed-size chunks with overlap."""

    def __init__(self, chunk_size: int = 500, overlap: int = 100):
        self.chunk_size = chunk_size
        self.overlap = overlap

    def chunk(self, text: str, doc_id: str) -> List[Document]:
        chunks = []
        start = 0
        chunk_num = 0

        while start < len(text):
            end = start + self.chunk_size
            chunk_text = text[start:end]

            # Try to break at sentence boundary
            if end < len(text):
                for sep in ['. ', '! ', '? ', '\n\n', '\n']:
                    last_sep = chunk_text.rfind(sep)
                    if last_sep > self.chunk_size // 2:
                        chunk_text = chunk_text[:last_sep + len(sep)]
                        end = start + len(chunk_text)
                        break

            chunks.append(Document(
                id=f"{doc_id}_chunk_{chunk_num}",
                content=chunk_text.strip(),
                metadata={
                    "parent_doc_id": doc_id,
                    "chunk_index": chunk_num,
                    "start_char": start,
                    "end_char": end,
                    "chunking": "fixed_size",
                }
            ))

            chunk_num += 1
            start = end - self.overlap
            if start >= end:
                break

        return chunks


class SentenceChunker(ChunkingStrategy):
    """Sentence-based chunking: groups sentences up to max_chunk_size."""

    def __init__(self, max_chunk_size: int = 500, overlap_sentences: int = 1):
        self.max_chunk_size = max_chunk_size
        self.overlap_sentences = overlap_sentences

    def _split_sentences(self, text: str) -> List[str]:
        return re.split(r'(?<=[.!?])\s+', text.strip())

    def chunk(self, text: str, doc_id: str) -> List[Document]:
        sentences = self._split_sentences(text)
        if not sentences:
            return [Document(id=f"{doc_id}_chunk_0", content=text, metadata={"parent_doc_id": doc_id, "chunking": "sentence"})]

        chunks = []
        current_sents: List[str] = []
        current_len = 0
        chunk_num = 0

        for sent in sentences:
            sent = sent.strip()
            if not sent:
                continue

            if current_len + len(sent) > self.max_chunk_size and current_sents:
                chunk_text = " ".join(current_sents)
                chunks.append(Document(
                    id=f"{doc_id}_chunk_{chunk_num}",
                    content=chunk_text,
                    metadata={"parent_doc_id": doc_id, "chunk_index": chunk_num, "chunking": "sentence"},
                ))
                chunk_num += 1
                # Keep overlap
                overlap = current_sents[-self.overlap_sentences:] if self.overlap_sentences else []
                current_sents = overlap
                current_len = sum(len(s) for s in current_sents)

            current_sents.append(sent)
            current_len += len(sent)

        if current_sents:
            chunks.append(Document(
                id=f"{doc_id}_chunk_{chunk_num}",
                content=" ".join(current_sents),
                metadata={"parent_doc_id": doc_id, "chunk_index": chunk_num, "chunking": "sentence"},
            ))

        return chunks


class ParagraphChunker(ChunkingStrategy):
    """Paragraph-based chunking: splits on double newlines."""

    def __init__(self, max_chunk_size: int = 1000):
        self.max_chunk_size = max_chunk_size

    def chunk(self, text: str, doc_id: str) -> List[Document]:
        paragraphs = [p.strip() for p in text.split('\n\n') if p.strip()]
        if not paragraphs:
            return [Document(id=f"{doc_id}_chunk_0", content=text, metadata={"parent_doc_id": doc_id, "chunking": "paragraph"})]

        chunks = []
        current_paras: List[str] = []
        current_len = 0
        chunk_num = 0

        for para in paragraphs:
            if current_len + len(para) > self.max_chunk_size and current_paras:
                chunks.append(Document(
                    id=f"{doc_id}_chunk_{chunk_num}",
                    content="\n\n".join(current_paras),
                    metadata={"parent_doc_id": doc_id, "chunk_index": chunk_num, "chunking": "paragraph"},
                ))
                chunk_num += 1
                current_paras = []
                current_len = 0

            current_paras.append(para)
            current_len += len(para)

        if current_paras:
            chunks.append(Document(
                id=f"{doc_id}_chunk_{chunk_num}",
                content="\n\n".join(current_paras),
                metadata={"parent_doc_id": doc_id, "chunk_index": chunk_num, "chunking": "paragraph"},
            ))

        return chunks


class RecursiveChunker(ChunkingStrategy):
    """Recursive character text splitter: tries multiple separators in order."""

    def __init__(self, chunk_size: int = 500, overlap: int = 100,
                 separators: Optional[List[str]] = None):
        self.chunk_size = chunk_size
        self.overlap = overlap
        self.separators = separators or ["\n\n", "\n", ". ", " ", ""]

    def _split(self, text: str, separators: List[str]) -> List[str]:
        if not text:
            return []

        if not separators:
            return [text]

        sep = separators[0]
        rest = separators[1:]

        if not sep:
            # Character-level split
            parts = [text[i:i + self.chunk_size] for i in range(0, len(text), self.chunk_size - self.overlap)]
            return parts

        splits = text.split(sep)
        result = []
        current = ""

        for piece in splits:
            candidate = (current + sep + piece) if current else piece
            if len(candidate) <= self.chunk_size:
                current = candidate
            else:
                if current:
                    result.append(current)
                if len(piece) > self.chunk_size:
                    result.extend(self._split(piece, rest))
                    current = ""
                else:
                    current = piece

        if current:
            result.append(current)

        return result

    def chunk(self, text: str, doc_id: str) -> List[Document]:
        pieces = self._split(text, self.separators)
        chunks = []
        for i, piece in enumerate(pieces):
            chunks.append(Document(
                id=f"{doc_id}_chunk_{i}",
                content=piece.strip(),
                metadata={"parent_doc_id": doc_id, "chunk_index": i, "chunking": "recursive"},
            ))
        return [c for c in chunks if c.content]


# ---------------------------------------------------------------------------
# Document Loaders
# ---------------------------------------------------------------------------

def load_txt(path: Path) -> List[Dict[str, Any]]:
    """Load plain text file."""
    content = path.read_text(encoding="utf-8")
    return [{"content": content, "id": path.stem, "metadata": {"source": str(path), "type": "txt"}}]


def load_markdown(path: Path) -> List[Dict[str, Any]]:
    """Load markdown file, preserving structure."""
    content = path.read_text(encoding="utf-8")
    return [{"content": content, "id": path.stem, "metadata": {"source": str(path), "type": "markdown"}}]


def load_csv(path: Path, text_columns: Optional[List[str]] = None) -> List[Dict[str, Any]]:
    """Load CSV/TSV file, treating each row as a document."""
    delimiter = '\t' if path.suffix.lower() in ('.tsv',) else ','
    documents = []
    with open(path, "r", encoding="utf-8", newline="") as f:
        reader = csv.DictReader(f, delimiter=delimiter)
        for i, row in enumerate(reader):
            if text_columns:
                content = " ".join(str(row.get(col, "")) for col in text_columns if row.get(col))
            else:
                content = " ".join(str(v) for v in row.values() if v)
            if content.strip():
                documents.append({
                    "content": content,
                    "id": f"{path.stem}_row_{i}",
                    "metadata": {"source": str(path), "type": "csv", "row": i, **row},
                })
    return documents


def load_json_file(path: Path) -> List[Dict[str, Any]]:
    """Load JSON file."""
    data = json.loads(path.read_text(encoding="utf-8"))
    documents = []
    if isinstance(data, list):
        for i, item in enumerate(data):
            if isinstance(item, dict):
                content = item.get("content") or item.get("text") or json.dumps(item)
                doc = {
                    "content": content,
                    "id": item.get("id") or f"{path.stem}_{i}",
                    "metadata": {k: v for k, v in item.items() if k not in ("content", "text", "id")},
                }
            else:
                doc = {"content": str(item), "id": f"{path.stem}_{i}", "metadata": {}}
            doc["metadata"]["source"] = str(path)
            doc["metadata"]["type"] = "json"
            documents.append(doc)
    elif isinstance(data, dict):
        documents.append({
            "content": data.get("content") or data.get("text") or json.dumps(data),
            "id": data.get("id") or path.stem,
            "metadata": {"source": str(path), "type": "json"},
        })
    return documents


def load_jsonl(path: Path) -> List[Dict[str, Any]]:
    """Load JSONL file."""
    documents = []
    with open(path, "r", encoding="utf-8") as f:
        for i, line in enumerate(f):
            line = line.strip()
            if not line:
                continue
            item = json.loads(line)
            content = item.get("content") or item.get("text") or item.get("question") or json.dumps(item)
            doc = {
                "content": content,
                "id": item.get("id") or f"{path.stem}_{i}",
                "metadata": {k: v for k, v in item.items() if k not in ("content", "text", "id")},
            }
            doc["metadata"]["source"] = str(path)
            doc["metadata"]["type"] = "jsonl"
            documents.append(doc)
    return documents


def load_pdf(path: Path) -> List[Dict[str, Any]]:
    """
    Load PDF file. Tries pymupdf (fitz), then pdfplumber, then falls back gracefully.
    """
    text_pages: List[str] = []

    # Try pymupdf first
    try:
        import fitz  # pymupdf
        doc = fitz.open(str(path))
        for page in doc:
            text_pages.append(page.get_text())
        doc.close()
    except ImportError:
        # Try pdfplumber
        try:
            import pdfplumber
            with pdfplumber.open(str(path)) as pdf:
                for page in pdf.pages:
                    t = page.extract_text()
                    if t:
                        text_pages.append(t)
        except ImportError:
            raise ImportError(
                "No PDF library available. Install one of:\n"
                "  pip install pymupdf    (recommended)\n"
                "  pip install pdfplumber"
            )

    content = "\n\n".join(text_pages)
    if not content.strip():
        return []

    return [{
        "content": content,
        "id": path.stem,
        "metadata": {"source": str(path), "type": "pdf", "pages": len(text_pages)},
    }]


# Map file extensions to loaders
DOCUMENT_LOADERS = {
    ".txt": load_txt,
    ".md": load_markdown,
    ".markdown": load_markdown,
    ".csv": load_csv,
    ".tsv": load_csv,
    ".json": load_json_file,
    ".jsonl": load_jsonl,
    ".pdf": load_pdf,
}


# ---------------------------------------------------------------------------
# BM25 (keyword-based scoring)
# ---------------------------------------------------------------------------

class SimpleBM25:
    """Simple BM25 scoring for hybrid search. No external dependencies."""

    def __init__(self, k1: float = 1.5, b: float = 0.75):
        self.k1 = k1
        self.b = b
        self._doc_freqs: Dict[str, int] = Counter()
        self._doc_lens: List[int] = []
        self._avg_dl: float = 0
        self._n_docs: int = 0
        self._tokenized_docs: List[List[str]] = []

    @staticmethod
    def _tokenize(text: str) -> List[str]:
        return re.findall(r'\b\w+\b', text.lower())

    def index(self, documents: List[str]):
        """Index a list of document strings."""
        self._tokenized_docs = [self._tokenize(d) for d in documents]
        self._n_docs = len(documents)
        self._doc_lens = [len(d) for d in self._tokenized_docs]
        self._avg_dl = sum(self._doc_lens) / max(self._n_docs, 1)

        self._doc_freqs = Counter()
        for tokens in self._tokenized_docs:
            seen = set(tokens)
            for t in seen:
                self._doc_freqs[t] += 1

    def score(self, query: str) -> np.ndarray:
        """Score all documents against a query. Returns array of scores."""
        query_tokens = self._tokenize(query)
        scores = np.zeros(self._n_docs)

        for token in query_tokens:
            df = self._doc_freqs.get(token, 0)
            if df == 0:
                continue
            idf = math.log((self._n_docs - df + 0.5) / (df + 0.5) + 1.0)

            for i, doc_tokens in enumerate(self._tokenized_docs):
                tf = doc_tokens.count(token)
                dl = self._doc_lens[i]
                numerator = tf * (self.k1 + 1)
                denominator = tf + self.k1 * (1 - self.b + self.b * dl / max(self._avg_dl, 1))
                scores[i] += idf * numerator / denominator

        return scores


# ---------------------------------------------------------------------------
# Retrieval Tool
# ---------------------------------------------------------------------------

class Retrieval(BaseTool):
    """
    Standard Retrieval (RAG) tool with embedding-based search.

    Features:
    - Embedding-based semantic search
    - Multiple embedding providers (Sentence Transformers, TF-IDF fallback)
    - Document loaders: txt, md, pdf, csv, tsv, json, jsonl
    - Chunking strategies: fixed-size, sentence, paragraph, recursive
    - Hybrid search: vector similarity + BM25 keyword matching
    - Persistent index storage
    - Metadata filtering
    - Score thresholds

    Usage:
        retrieval = Retrieval()
        retrieval.add_documents([{"content": "...", "id": "doc1"}])
        result = await retrieval.execute(query="your question", top_k=5)
    """

    CHUNKING_STRATEGIES = {
        "fixed": FixedSizeChunker,
        "sentence": SentenceChunker,
        "paragraph": ParagraphChunker,
        "recursive": RecursiveChunker,
    }

    def __init__(
        self,
        embedding_provider: Optional[EmbeddingProvider] = None,
        chunk_size: int = 500,
        chunk_overlap: int = 100,
        chunking_strategy: str = "fixed",
        index_path: Optional[str] = None,
        knowledge_base_path: Optional[str] = None,
        enable_hybrid_search: bool = True,
        hybrid_alpha: float = 0.7,
    ):
        """
        Initialize the retrieval tool.

        Args:
            embedding_provider: Provider for text embeddings
            chunk_size: Size of document chunks in characters
            chunk_overlap: Overlap between chunks in characters
            chunking_strategy: One of 'fixed', 'sentence', 'paragraph', 'recursive'
            index_path: Path to persist the index
            knowledge_base_path: Path to a directory to auto-load documents from
            enable_hybrid_search: Combine vector + BM25 search (default True)
            hybrid_alpha: Weight for vector score in hybrid (0=BM25 only, 1=vector only)
        """
        super().__init__(
            metadata=ToolMetadata(
                name="retrieval",
                description=(
                    "Search a knowledge base using semantic similarity. "
                    "Returns the most relevant documents/passages for a given query. "
                    "Use this tool when you need to find information from the indexed knowledge base."
                ),
                category=ToolCategory.INFORMATION_RETRIEVAL,
                parameters=[
                    ParameterSpec(
                        name="query",
                        type=ParameterType.STRING,
                        description="The search query to find relevant documents",
                        required=True,
                        min_length=1,
                        max_length=1000,
                    ),
                    ParameterSpec(
                        name="top_k",
                        type=ParameterType.INTEGER,
                        description="Number of top results to return",
                        required=False,
                        default=5,
                        min_value=1,
                        max_value=50,
                    ),
                    ParameterSpec(
                        name="score_threshold",
                        type=ParameterType.FLOAT,
                        description="Minimum similarity score threshold (0-1)",
                        required=False,
                        default=0.0,
                        min_value=0.0,
                        max_value=1.0,
                    ),
                    ParameterSpec(
                        name="filter_metadata",
                        type=ParameterType.OBJECT,
                        description="Filter results by metadata fields",
                        required=False,
                    ),
                ],
                returns={
                    "type": "object",
                    "properties": {
                        "results": {
                            "type": "array",
                            "items": {
                                "type": "object",
                                "properties": {
                                    "content": {"type": "string"},
                                    "score": {"type": "number"},
                                    "metadata": {"type": "object"},
                                },
                            },
                        },
                        "total_found": {"type": "integer"},
                    },
                },
                timeout_seconds=30,
                tags=["retrieval", "rag", "search", "knowledge-base", "semantic"],
                examples=[
                    {
                        "query": "What is machine learning?",
                        "top_k": 3,
                        "output": {
                            "results": [
                                {
                                    "content": "Machine learning is a subset of AI...",
                                    "score": 0.92,
                                    "metadata": {"source": "ml_intro.txt"},
                                }
                            ],
                            "total_found": 3,
                        },
                    }
                ],
            )
        )

        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.index_path = index_path
        self.enable_hybrid_search = enable_hybrid_search
        self.hybrid_alpha = hybrid_alpha

        # Initialize chunker
        if chunking_strategy in self.CHUNKING_STRATEGIES:
            if chunking_strategy == "fixed":
                self._chunker = FixedSizeChunker(chunk_size, chunk_overlap)
            elif chunking_strategy == "sentence":
                self._chunker = SentenceChunker(chunk_size)
            elif chunking_strategy == "paragraph":
                self._chunker = ParagraphChunker(chunk_size)
            elif chunking_strategy == "recursive":
                self._chunker = RecursiveChunker(chunk_size, chunk_overlap)
        else:
            self._chunker = FixedSizeChunker(chunk_size, chunk_overlap)

        # Initialize embedding provider
        if embedding_provider is not None:
            self.embedding_provider = embedding_provider
        else:
            try:
                import importlib.util
                if importlib.util.find_spec("sentence_transformers") is not None:
                    self.embedding_provider = SentenceTransformerEmbedding()
                else:
                    raise ImportError("sentence_transformers not found")
            except (ImportError, Exception):
                logger.warning(
                    "\u26a0\ufe0f  sentence-transformers not installed. Using simple TF-IDF embeddings.\n"
                    "    For better results: pip install sentence-transformers"
                )
                self.embedding_provider = SimpleEmbedding()

        # Document storage
        self.documents: Dict[str, Document] = {}
        self.embeddings_matrix: Optional[np.ndarray] = None
        self.doc_ids: List[str] = []

        # BM25 index for hybrid search
        self._bm25: Optional[SimpleBM25] = None

        # Load existing index if available
        if index_path and os.path.exists(index_path):
            self._load_index(index_path)

        # Auto-load knowledge base directory
        if knowledge_base_path and os.path.isdir(knowledge_base_path):
            self._load_directory(knowledge_base_path)

    def _load_directory(self, dir_path: str):
        """Load all supported files from a directory."""
        loaded = 0
        p = Path(dir_path)
        for ext, loader_fn in DOCUMENT_LOADERS.items():
            for file_path in p.glob(f"*{ext}"):
                try:
                    docs = loader_fn(file_path)
                    if docs:
                        loaded += self.add_documents(docs)
                except Exception as e:
                    logger.warning(f"Failed to load {file_path}: {e}")
        if loaded:
            logger.info(f"Loaded {loaded} documents/chunks from {dir_path}")

    async def initialize(self) -> None:
        """Initialize the retrieval tool."""
        await super().initialize()
        logger.info("Retrieval tool initialized")

    def _chunk_text(self, text: str, doc_id: str) -> List[Document]:
        """Split text into chunks using the configured strategy."""
        return self._chunker.chunk(text, doc_id)

    def add_documents(
        self,
        documents: List[Dict[str, Any]],
        chunk: bool = True,
    ) -> int:
        """
        Add documents to the index.

        Args:
            documents: List of documents with 'content', optional 'id' and 'metadata'
            chunk: Whether to chunk documents

        Returns:
            Number of documents/chunks added
        """
        all_chunks = []

        for doc in documents:
            content = doc.get("content", "")
            if not content:
                continue

            doc_id = doc.get("id") or hashlib.md5(content.encode()).hexdigest()[:12]
            metadata = doc.get("metadata", {})

            if chunk and len(content) > self.chunk_size:
                chunks = self._chunk_text(content, doc_id)
                for c in chunks:
                    c.metadata.update(metadata)
                all_chunks.extend(chunks)
            else:
                all_chunks.append(Document(
                    id=doc_id,
                    content=content,
                    metadata=metadata,
                ))

        if not all_chunks:
            return 0

        # Generate embeddings
        texts = [c.content for c in all_chunks]

        if isinstance(self.embedding_provider, SimpleEmbedding):
            all_texts = [d.content for d in self.documents.values()] + texts
            self.embedding_provider.fit(all_texts)

        embeddings = self.embedding_provider.embed(texts)

        for i, chunk in enumerate(all_chunks):
            chunk.embedding = embeddings[i]
            self.documents[chunk.id] = chunk
            self.doc_ids.append(chunk.id)

        self._rebuild_embedding_matrix()
        self._rebuild_bm25_index()

        logger.info(f"Added {len(all_chunks)} documents/chunks to index")
        return len(all_chunks)

    def add_from_file(
        self,
        file_path: str,
        file_type: str = "auto",
        chunk: bool = True,
    ) -> int:
        """
        Add documents from a file.

        Supports: txt, md, pdf, csv, tsv, json, jsonl

        Args:
            file_path: Path to the file
            file_type: File type (auto-detected from extension by default)
            chunk: Whether to chunk documents

        Returns:
            Number of documents added
        """
        path = Path(file_path)
        if not path.exists():
            raise FileNotFoundError(f"File not found: {file_path}")

        # Auto-detect file type
        ext = path.suffix.lower()
        if file_type != "auto":
            ext = f".{file_type}" if not file_type.startswith(".") else file_type

        loader_fn = DOCUMENT_LOADERS.get(ext)
        if loader_fn is None:
            # Fallback to plain text
            loader_fn = load_txt

        documents = loader_fn(path)
        return self.add_documents(documents, chunk=chunk)

    def _rebuild_embedding_matrix(self):
        """Rebuild the embedding matrix from stored documents."""
        if not self.documents:
            self.embeddings_matrix = None
            self.doc_ids = []
            return

        self.doc_ids = list(self.documents.keys())
        embeddings = [self.documents[doc_id].embedding for doc_id in self.doc_ids]
        self.embeddings_matrix = np.vstack(embeddings)

    def _rebuild_bm25_index(self):
        """Rebuild the BM25 index for hybrid search."""
        if not self.enable_hybrid_search or not self.documents:
            self._bm25 = None
            return
        self._bm25 = SimpleBM25()
        texts = [self.documents[doc_id].content for doc_id in self.doc_ids]
        self._bm25.index(texts)

    def _cosine_similarity(self, query_embedding: np.ndarray) -> np.ndarray:
        """Compute cosine similarity between query and all documents."""
        if self.embeddings_matrix is None:
            return np.array([])

        query_norm = query_embedding / (np.linalg.norm(query_embedding) + 1e-8)
        doc_norms = self.embeddings_matrix / (
            np.linalg.norm(self.embeddings_matrix, axis=1, keepdims=True) + 1e-8
        )
        similarities = np.dot(doc_norms, query_norm)
        return similarities

    async def _execute(
        self,
        query: str,
        top_k: int = 5,
        score_threshold: float = 0.0,
        filter_metadata: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> Dict[str, Any]:
        """
        Execute retrieval query.

        Uses hybrid search (vector + BM25) by default if enabled.
        """
        if not self.documents:
            return {
                "results": [],
                "total_found": 0,
                "message": "No documents in index. Add documents first using add_documents() or add_from_file().",
            }

        # Get query embedding
        query_embedding = self.embedding_provider.embed_query(query)

        # Vector similarity scores
        vector_scores = self._cosine_similarity(query_embedding)

        # Hybrid: combine with BM25
        if self.enable_hybrid_search and self._bm25 is not None:
            bm25_scores = self._bm25.score(query)
            # Normalize BM25 scores to [0, 1]
            bm25_max = bm25_scores.max() if bm25_scores.max() > 0 else 1.0
            bm25_norm = bm25_scores / bm25_max
            # Normalize vector scores to [0, 1]
            vec_max = vector_scores.max() if vector_scores.max() > 0 else 1.0
            vec_norm = vector_scores / vec_max
            combined_scores = self.hybrid_alpha * vec_norm + (1 - self.hybrid_alpha) * bm25_norm
        else:
            combined_scores = vector_scores

        top_indices = np.argsort(combined_scores)[::-1]

        results = []
        for idx in top_indices:
            if len(results) >= top_k:
                break

            doc_id = self.doc_ids[idx]
            doc = self.documents[doc_id]
            score = float(combined_scores[idx])

            if score < score_threshold:
                continue

            if filter_metadata:
                match = True
                for key, value in filter_metadata.items():
                    if doc.metadata.get(key) != value:
                        match = False
                        break
                if not match:
                    continue

            results.append({
                "content": doc.content,
                "score": round(score, 4),
                "id": doc.id,
                "metadata": doc.metadata,
                "rank": len(results) + 1,
            })

        return {
            "results": results,
            "total_found": len(results),
            "query": query,
        }

    def save_index(self, path: Optional[str] = None):
        """Save the index to disk."""
        path = path or self.index_path
        if not path:
            raise ValueError("No index path specified")

        data = {
            "documents": {
                doc_id: {
                    "id": doc.id,
                    "content": doc.content,
                    "metadata": doc.metadata,
                    "embedding": doc.embedding.tolist() if doc.embedding is not None else None,
                }
                for doc_id, doc in self.documents.items()
            },
            "doc_ids": self.doc_ids,
            "chunk_size": self.chunk_size,
            "chunk_overlap": self.chunk_overlap,
        }

        with open(path, "wb") as f:
            pickle.dump(data, f)

        logger.info(f"Saved index to {path}")

    def _load_index(self, path: str):
        """Load the index from disk."""
        with open(path, "rb") as f:
            data = pickle.load(f)

        self.documents = {}
        for doc_id, doc_data in data["documents"].items():
            self.documents[doc_id] = Document(
                id=doc_data["id"],
                content=doc_data["content"],
                metadata=doc_data["metadata"],
                embedding=np.array(doc_data["embedding"]) if doc_data["embedding"] else None,
            )

        self.doc_ids = data["doc_ids"]
        self.chunk_size = data.get("chunk_size", self.chunk_size)
        self.chunk_overlap = data.get("chunk_overlap", self.chunk_overlap)

        self._rebuild_embedding_matrix()
        self._rebuild_bm25_index()
        logger.info(f"Loaded index from {path} with {len(self.documents)} documents")

    def clear(self):
        """Clear all documents from the index."""
        self.documents = {}
        self.embeddings_matrix = None
        self.doc_ids = []
        self._bm25 = None
        logger.info("Cleared retrieval index")

    @property
    def num_documents(self) -> int:
        """Return the number of documents in the index."""
        return len(self.documents)
