"""BM25 indexing implementation for Agent learning memory.

This module provides BM25 (Best Matching 25) indexing for efficient
text-based retrieval of agent memories. BM25 is a probabilistic retrieval
framework that ranks documents based on query terms.
"""

import re
import math
import pickle
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
from dataclasses import dataclass, field
from datetime import datetime


@dataclass
class MemoryDocument:
    """A document stored in the memory index.

    Attributes:
        doc_id: Unique document identifier
        content: Text content for indexing
        memory_type: Type of memory (trade, market, decision, error)
        timestamp: When the memory was created
        metadata: Additional structured data
        importance: Importance score (0.0 to 1.0)
        access_count: Number of times this memory was retrieved
        last_accessed: Last access timestamp
    """

    doc_id: str
    content: str
    memory_type: str
    timestamp: datetime = field(default_factory=datetime.now)
    metadata: Dict[str, Any] = field(default_factory=dict)
    importance: float = 0.5
    access_count: int = 0
    last_accessed: Optional[datetime] = None

    def to_dict(self) -> Dict[str, Any]:
        """Convert document to dictionary for serialization."""
        return {
            "doc_id": self.doc_id,
            "content": self.content,
            "memory_type": self.memory_type,
            "timestamp": self.timestamp.isoformat(),
            "metadata": self.metadata,
            "importance": self.importance,
            "access_count": self.access_count,
            "last_accessed": self.last_accessed.isoformat() if self.last_accessed else None,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "MemoryDocument":
        """Create document from dictionary."""
        return cls(
            doc_id=data["doc_id"],
            content=data["content"],
            memory_type=data["memory_type"],
            timestamp=datetime.fromisoformat(data["timestamp"]),
            metadata=data.get("metadata", {}),
            importance=data.get("importance", 0.5),
            access_count=data.get("access_count", 0),
            last_accessed=datetime.fromisoformat(data["last_accessed"])
            if data.get("last_accessed")
            else None,
        )


class BM25Index:
    """BM25 index for efficient text retrieval.

    Implements the BM25 algorithm for ranking documents by relevance
    to a query. Supports incremental updates and persistence.

    BM25 formula:
        score(q, d) = sum(IDF(q_i) * (f(q_i, d) * (k1 + 1)) /
                              (f(q_i, d) + k1 * (1 - b + b * |d| / avgdl)))

    Where:
        - q_i: query term i
        - f(q_i, d): frequency of term q_i in document d
        - |d|: length of document d
        - avgdl: average document length
        - k1: term frequency saturation parameter (default: 1.5)
        - b: length normalization parameter (default: 0.75)
    """

    def __init__(
        self,
        k1: float = 1.5,
        b: float = 0.75,
        delta: float = 1.0,
        min_doc_freq: int = 1,
    ):
        """Initialize BM25 index.

        Args:
            k1: Term frequency saturation parameter
            b: Length normalization parameter (0-1)
            delta: BM25+ delta parameter for document length normalization
            min_doc_freq: Minimum document frequency for a term to be indexed
        """
        self.k1 = k1
        self.b = b
        self.delta = delta
        self.min_doc_freq = min_doc_freq

        # Document storage
        self.documents: Dict[str, MemoryDocument] = {}

        # Inverted index: term -> {doc_id: term_frequency}
        self.inverted_index: Dict[str, Dict[str, int]] = {}

        # Document lengths and statistics
        self.doc_lengths: Dict[str, int] = {}
        self.total_doc_length: int = 0
        self.avg_doc_length: float = 0.0
        self.num_docs: int = 0

        # Term document frequencies
        self.doc_freqs: Dict[str, int] = {}

        # IDF cache
        self._idf_cache: Dict[str, float] = {}

    def _tokenize(self, text: str) -> List[str]:
        """Tokenize text into terms.

        Args:
            text: Input text to tokenize

        Returns:
            List of normalized tokens
        """
        # Convert to lowercase and extract alphanumeric tokens
        text = text.lower()
        tokens = re.findall(r"\b[a-z0-9]+\b", text)
        return tokens

    def _compute_idf(self, term: str) -> float:
        """Compute IDF (Inverse Document Frequency) for a term.

        Uses BM25's modified IDF formula:
            IDF(q_i) = log((N - n(q_i) + 0.5) / (n(q_i) + 0.5))

        Args:
            term: The term to compute IDF for

        Returns:
            IDF score for the term
        """
        if term in self._idf_cache:
            return self._idf_cache[term]

        n_q = self.doc_freqs.get(term, 0)
        if n_q == 0:
            return 0.0

        # BM25 IDF formula
        idf = math.log((self.num_docs - n_q + 0.5) / (n_q + 0.5) + 1.0)
        self._idf_cache[term] = idf
        return idf

    def _compute_score(self, query_terms: List[str], doc_id: str) -> float:
        """Compute BM25 score for a document given query terms.

        Args:
            query_terms: Tokenized query terms
            doc_id: Document ID to score

        Returns:
            BM25 relevance score
        """
        if doc_id not in self.documents:
            return 0.0

        score = 0.0
        doc_length = self.doc_lengths[doc_id]

        # Document length normalization
        norm_factor = 1 - self.b + self.b * (doc_length / self.avg_doc_length) if self.avg_doc_length > 0 else 1

        for term in query_terms:
            if term not in self.inverted_index:
                continue

            # Term frequency in document
            tf = self.inverted_index[term].get(doc_id, 0)
            if tf == 0:
                continue

            # IDF for term
            idf = self._compute_idf(term)

            # BM25 term score
            numerator = tf * (self.k1 + 1)
            denominator = tf + self.k1 * norm_factor
            term_score = idf * (numerator / denominator + self.delta)

            score += term_score

        return score

    def add_document(self, doc: MemoryDocument) -> None:
        """Add a document to the index.

        Args:
            doc: MemoryDocument to add
        """
        doc_id = doc.doc_id

        # Remove existing document if present (for updates)
        if doc_id in self.documents:
            self.remove_document(doc_id)

        # Store document
        self.documents[doc_id] = doc

        # Tokenize and index
        tokens = self._tokenize(doc.content)
        token_counts: Dict[str, int] = {}

        for token in tokens:
            token_counts[token] = token_counts.get(token, 0) + 1

        # Update inverted index
        for token, count in token_counts.items():
            if token not in self.inverted_index:
                self.inverted_index[token] = {}
                self.doc_freqs[token] = 0

            self.inverted_index[token][doc_id] = count
            self.doc_freqs[token] += 1

        # Update document statistics
        doc_length = len(tokens)
        self.doc_lengths[doc_id] = doc_length
        self.total_doc_length += doc_length
        self.num_docs += 1
        self.avg_doc_length = self.total_doc_length / self.num_docs if self.num_docs > 0 else 0.0

        # Clear IDF cache
        self._idf_cache.clear()

    def remove_document(self, doc_id: str) -> bool:
        """Remove a document from the index.

        Args:
            doc_id: Document ID to remove

        Returns:
            True if document was removed, False if not found
        """
        if doc_id not in self.documents:
            return False

        # Update inverted index
        tokens_to_remove: List[str] = []
        for term, postings in self.inverted_index.items():
            if doc_id in postings:
                del postings[doc_id]
                self.doc_freqs[term] -= 1
                if self.doc_freqs[term] <= 0:
                    tokens_to_remove.append(term)

        # Clean up empty terms
        for term in tokens_to_remove:
            del self.inverted_index[term]
            del self.doc_freqs[term]

        # Update statistics
        doc_length = self.doc_lengths.get(doc_id, 0)
        self.total_doc_length -= doc_length
        self.num_docs -= 1
        self.avg_doc_length = self.total_doc_length / self.num_docs if self.num_docs > 0 else 0.0

        # Remove document
        del self.documents[doc_id]
        del self.doc_lengths[doc_id]

        # Clear IDF cache
        self._idf_cache.clear()

        return True

    def search(
        self,
        query: str,
        top_k: int = 10,
        memory_type: Optional[str] = None,
        min_score: float = 0.0,
    ) -> List[Tuple[MemoryDocument, float]]:
        """Search for documents matching the query.

        Args:
            query: Search query text
            top_k: Maximum number of results to return
            memory_type: Filter by memory type (optional)
            min_score: Minimum score threshold

        Returns:
            List of (document, score) tuples sorted by relevance
        """
        query_terms = self._tokenize(query)
        if not query_terms:
            return []

        # Score all matching documents
        candidates: Dict[str, float] = {}

        # Get candidate documents from query terms
        candidate_ids: set = set()
        for term in query_terms:
            if term in self.inverted_index:
                candidate_ids.update(self.inverted_index[term].keys())

        # Filter by memory type if specified
        if memory_type:
            candidate_ids = {
                doc_id for doc_id in candidate_ids
                if self.documents[doc_id].memory_type == memory_type
            }

        # Score candidates
        for doc_id in candidate_ids:
            score = self._compute_score(query_terms, doc_id)
            if score >= min_score:
                candidates[doc_id] = score

        # Sort by score and return top_k
        sorted_results = sorted(candidates.items(), key=lambda x: x[1], reverse=True)
        top_results = sorted_results[:top_k]

        # Update access statistics
        results: List[Tuple[MemoryDocument, float]] = []
        for doc_id, score in top_results:
            doc = self.documents[doc_id]
            doc.access_count += 1
            doc.last_accessed = datetime.now()
            results.append((doc, score))

        return results

    def get_document(self, doc_id: str) -> Optional[MemoryDocument]:
        """Get a document by ID.

        Args:
            doc_id: Document ID

        Returns:
            MemoryDocument if found, None otherwise
        """
        return self.documents.get(doc_id)

    def update_document(self, doc_id: str, **updates: Any) -> bool:
        """Update document fields.

        Args:
            doc_id: Document ID to update
            **updates: Fields to update

        Returns:
            True if updated, False if not found
        """
        if doc_id not in self.documents:
            return False

        doc = self.documents[doc_id]

        # Update fields
        if "content" in updates:
            # Re-index if content changes
            doc.content = updates["content"]
            self.remove_document(doc_id)
            self.add_document(doc)
        if "importance" in updates:
            doc.importance = updates["importance"]
        if "metadata" in updates:
            doc.metadata.update(updates["metadata"])

        return True

    def get_stats(self) -> Dict[str, Any]:
        """Get index statistics.

        Returns:
            Dictionary with index statistics
        """
        return {
            "num_documents": self.num_docs,
            "num_terms": len(self.inverted_index),
            "avg_doc_length": self.avg_doc_length,
            "total_doc_length": self.total_doc_length,
            "memory_types": {
                mem_type: sum(1 for d in self.documents.values() if d.memory_type == mem_type)
                for mem_type in set(d.memory_type for d in self.documents.values())
            },
        }

    def save(self, path: Path) -> None:
        """Save index to disk using pickle.

        Note: Pickle is used for local index storage. This is acceptable
        since the data is generated and consumed by the same system.

        Args:
            path: Path to save index
        """
        data = {
            "k1": self.k1,
            "b": self.b,
            "delta": self.delta,
            "min_doc_freq": self.min_doc_freq,
            "documents": {k: v.to_dict() for k, v in self.documents.items()},
            "inverted_index": dict(self.inverted_index),
            "doc_lengths": dict(self.doc_lengths),
            "total_doc_length": self.total_doc_length,
            "avg_doc_length": self.avg_doc_length,
            "num_docs": self.num_docs,
            "doc_freqs": dict(self.doc_freqs),
        }
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "wb") as f:
            pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)

    def load(self, path: Path) -> bool:
        """Load index from disk.

        Args:
            path: Path to load index from

        Returns:
            True if loaded successfully, False otherwise
        """
        if not path.exists():
            return False

        try:
            with open(path, "rb") as f:
                data = pickle.load(f)

            self.k1 = data.get("k1", 1.5)
            self.b = data.get("b", 0.75)
            self.delta = data.get("delta", 1.0)
            self.min_doc_freq = data.get("min_doc_freq", 1)

            self.documents = {
                k: MemoryDocument.from_dict(v) for k, v in data["documents"].items()
            }
            self.inverted_index = data["inverted_index"]
            self.doc_lengths = data["doc_lengths"]
            self.total_doc_length = data["total_doc_length"]
            self.avg_doc_length = data["avg_doc_length"]
            self.num_docs = data["num_docs"]
            self.doc_freqs = data["doc_freqs"]
            self._idf_cache.clear()

            return True
        except (pickle.PickleError, KeyError, IOError):
            return False
