"""
BM25 Index - Inverted index stored in SQLite for efficient keyword search

Implements the BM25 (Best Matching 25) algorithm for ranking documents
by relevance to a search query. All data is stored in SQLite for persistence.

BM25 formula:
score(D, Q) = Σ IDF(qi) * (f(qi, D) * (k1 + 1)) / (f(qi, D) + k1 * (1 - b + b * |D|/avgdl))

Where:
- f(qi, D) = term frequency of qi in document D
- |D| = document length
- avgdl = average document length
- k1 = term frequency saturation parameter (typically 1.2-2.0)
- b = length normalization parameter (typically 0.75)
- IDF(qi) = log((N - n(qi) + 0.5) / (n(qi) + 0.5))
- N = total number of documents
- n(qi) = number of documents containing qi
"""

import math
import re
from dataclasses import dataclass
from typing import Optional

from groknroll.storage.database import Database
from groknroll.storage.models import CodeSummary


@dataclass
class BM25Result:
    """Result from BM25 search"""

    summary_id: int
    score: float
    scope_type: str
    scope_path: str
    scope_name: Optional[str]
    summary: Optional[str]
    keywords: Optional[list[str]]


class BM25Index:
    """
    BM25 inverted index stored in SQLite.

    Indexes code summaries for efficient keyword-based retrieval.
    Supports field-weighted scoring (name, summary, docstring, keywords).
    """

    # BM25 parameters
    K1 = 1.5  # Term frequency saturation
    B = 0.75  # Length normalization

    # Field weights for scoring
    FIELD_WEIGHTS = {
        "name": 3.0,      # High weight for name matches
        "keywords": 2.5,  # High weight for keyword matches
        "summary": 1.5,   # Medium weight for summary matches
        "docstring": 1.0, # Base weight for docstring matches
    }

    # Stop words to filter out
    STOP_WORDS = {
        "the", "a", "an", "is", "are", "was", "were", "be", "been", "being",
        "have", "has", "had", "do", "does", "did", "will", "would", "could",
        "should", "may", "might", "must", "shall", "can", "need", "dare",
        "ought", "used", "to", "of", "in", "for", "on", "with", "at", "by",
        "from", "as", "into", "through", "during", "before", "after", "above",
        "below", "between", "under", "again", "further", "then", "once", "here",
        "there", "when", "where", "why", "how", "all", "each", "few", "more",
        "most", "other", "some", "such", "no", "nor", "not", "only", "own",
        "same", "so", "than", "too", "very", "just", "and", "but", "if", "or",
        "because", "until", "while", "this", "that", "these", "those", "it",
        "its", "self", "def", "class", "return", "none", "true", "false",
        "import", "from", "as", "try", "except", "finally", "raise", "pass",
    }

    def __init__(
        self,
        database: Database,
        project_id: int,
        verbose: bool = False,
    ):
        """
        Initialize BM25 index.

        Args:
            database: Database instance
            project_id: Project ID
            verbose: Whether to print progress
        """
        self.database = database
        self.project_id = project_id
        self.verbose = verbose

        # Cache for BM25 parameters
        self._total_docs: Optional[int] = None
        self._avg_doc_length: Optional[float] = None

    def index_summaries(self, force_rebuild: bool = False) -> int:
        """
        Index all summaries for the project.

        Args:
            force_rebuild: If True, clear and rebuild entire index

        Returns:
            Number of documents indexed
        """
        if force_rebuild:
            self.database.clear_bm25_index(self.project_id)

        # Get all summaries
        summaries = []
        for scope_type in ["module", "class", "function", "package"]:
            summaries.extend(
                self.database.get_summaries_by_type(self.project_id, scope_type)
            )

        if self.verbose:
            print(f"🔍 Indexing {len(summaries)} summaries for BM25...")

        indexed_count = 0
        for summary in summaries:
            self.index_document(summary)
            indexed_count += 1

        # Update document frequencies
        self._update_document_frequencies()

        # Invalidate cache
        self._total_docs = None
        self._avg_doc_length = None

        if self.verbose:
            print(f"✅ Indexed {indexed_count} documents")

        return indexed_count

    def index_document(self, summary: CodeSummary) -> None:
        """
        Index a single document (summary).

        Args:
            summary: CodeSummary to index
        """
        # Clear existing occurrences for this summary
        self.database.delete_term_occurrences_for_summary(summary.id)

        # Collect terms from different fields
        fields = {
            "name": summary.scope_name or "",
            "summary": summary.summary or "",
            "docstring": summary.docstring or "",
            "keywords": " ".join(summary.keywords or []),
        }

        # Index each field
        for field_type, text in fields.items():
            terms = self._tokenize(text)
            term_freqs = self._count_terms(terms)

            for term, freq in term_freqs.items():
                code_term = self.database.get_or_create_term(self.project_id, term)
                self.database.add_term_occurrence(
                    term_id=code_term.id,
                    summary_id=summary.id,
                    term_freq=freq,
                    field_type=field_type,
                )

    def search(
        self,
        query: str,
        limit: int = 20,
        scope_types: Optional[list[str]] = None,
    ) -> list[BM25Result]:
        """
        Search using BM25 ranking.

        Args:
            query: Search query
            limit: Maximum results to return
            scope_types: Filter by scope types (optional)

        Returns:
            List of BM25Result sorted by score descending
        """
        # Tokenize query
        query_terms = self._tokenize(query)
        if not query_terms:
            return []

        # Get BM25 parameters
        total_docs = self._get_total_docs()
        avg_doc_length = self._get_avg_doc_length()

        if total_docs == 0:
            return []

        # Search for terms
        term_data = self.database.search_terms(self.project_id, query_terms)

        # Calculate scores per document
        doc_scores: dict[int, float] = {}
        doc_info: dict[int, tuple[str, str, Optional[str], Optional[str], Optional[list[str]]]] = {}

        for code_term, occurrences in term_data:
            # IDF calculation
            idf = self._calculate_idf(code_term.document_freq, total_docs)

            for occ in occurrences:
                summary_id = occ.summary_id

                # Get document length (approximate from summary)
                if summary_id not in doc_info:
                    summary = self.database.get_summary_by_id(summary_id)
                    if summary:
                        if scope_types and summary.scope_type not in scope_types:
                            continue
                        doc_length = len(
                            (summary.summary or "") + " " + (summary.docstring or "")
                        )
                        doc_info[summary_id] = (
                            summary.scope_type,
                            summary.scope_path,
                            summary.scope_name,
                            summary.summary,
                            summary.keywords,
                        )
                    else:
                        continue
                else:
                    # Already cached
                    pass

                # Get field weight
                field_weight = self.FIELD_WEIGHTS.get(occ.field_type, 1.0)

                # Calculate BM25 score component
                tf = occ.term_freq
                doc_length = len(
                    (doc_info[summary_id][3] or "") + " " +
                    (self.database.get_summary_by_id(summary_id).docstring or "")
                ) if summary_id in doc_info else avg_doc_length

                score = self._calculate_bm25_score(
                    tf, doc_length, avg_doc_length, idf
                ) * field_weight

                if summary_id not in doc_scores:
                    doc_scores[summary_id] = 0.0
                doc_scores[summary_id] += score

        # Sort by score and limit
        sorted_results = sorted(doc_scores.items(), key=lambda x: -x[1])[:limit]

        # Build result objects
        results = []
        for summary_id, score in sorted_results:
            if summary_id in doc_info:
                scope_type, scope_path, scope_name, summary, keywords = doc_info[summary_id]
                results.append(BM25Result(
                    summary_id=summary_id,
                    score=score,
                    scope_type=scope_type,
                    scope_path=scope_path,
                    scope_name=scope_name,
                    summary=summary,
                    keywords=keywords,
                ))

        return results

    def _tokenize(self, text: str) -> list[str]:
        """
        Tokenize text into terms for indexing/searching.

        Args:
            text: Text to tokenize

        Returns:
            List of lowercase terms
        """
        if not text:
            return []

        # Split on non-alphanumeric, convert to lowercase
        words = re.findall(r"[a-zA-Z0-9]+", text.lower())

        # Split camelCase and snake_case
        expanded = []
        for word in words:
            # Split camelCase
            parts = re.findall(r"[a-z]+|[A-Z][a-z]*|[0-9]+", word)
            expanded.extend([p.lower() for p in parts if p])

        # Filter stop words and short words
        return [w for w in expanded if w not in self.STOP_WORDS and len(w) > 1]

    def _count_terms(self, terms: list[str]) -> dict[str, int]:
        """Count term frequencies."""
        counts: dict[str, int] = {}
        for term in terms:
            counts[term] = counts.get(term, 0) + 1
        return counts

    def _calculate_idf(self, doc_freq: int, total_docs: int) -> float:
        """
        Calculate Inverse Document Frequency.

        IDF(qi) = log((N - n(qi) + 0.5) / (n(qi) + 0.5) + 1)
        """
        if doc_freq == 0:
            return 0.0
        return math.log(
            (total_docs - doc_freq + 0.5) / (doc_freq + 0.5) + 1
        )

    def _calculate_bm25_score(
        self,
        term_freq: int,
        doc_length: float,
        avg_doc_length: float,
        idf: float,
    ) -> float:
        """
        Calculate BM25 score component for a single term.

        score = IDF * (f * (k1 + 1)) / (f + k1 * (1 - b + b * |D|/avgdl))
        """
        if avg_doc_length == 0:
            return 0.0

        numerator = term_freq * (self.K1 + 1)
        denominator = term_freq + self.K1 * (
            1 - self.B + self.B * (doc_length / avg_doc_length)
        )

        return idf * (numerator / denominator) if denominator > 0 else 0.0

    def _update_document_frequencies(self) -> None:
        """Update document frequencies for all terms."""
        # This is handled by the database when terms are added
        # Here we could optimize by batch updating
        pass

    def _get_total_docs(self) -> int:
        """Get total number of indexed documents."""
        if self._total_docs is None:
            self._total_docs = self.database.get_total_documents(self.project_id)
        return self._total_docs

    def _get_avg_doc_length(self) -> float:
        """Get average document length."""
        if self._avg_doc_length is None:
            self._avg_doc_length = self.database.get_average_document_length(
                self.project_id
            )
        return self._avg_doc_length

    def clear(self) -> None:
        """Clear the entire BM25 index for this project."""
        self.database.clear_bm25_index(self.project_id)
        self._total_docs = None
        self._avg_doc_length = None
