"""
Keyword-based search for modules and functions.

Provides semantic search capabilities by matching query keywords
against extracted keywords in the index using BM25 ranking.

Identifier names (function/module names) are given much higher weight than keywords.

Author: Cursor(Auto)
"""

import fnmatch
from typing import Any

from rank_bm25 import BM25Okapi

from cicada.utils import split_identifier


class KeywordSearcher:
    """Search for modules and functions by keywords using BM25 ranking."""

    # Boost multiplier for identifier name matches
    # When query keyword matches the function/module name, multiply the score by this
    IDENTIFIER_MATCH_BOOST = 10.0

    def __init__(self, index: dict[str, Any]):
        """
        Initialize the keyword searcher.

        Args:
            index: The Cicada index dictionary containing modules and metadata
        """
        self.index = index
        self.bm25, self.document_map = self._initialize_bm25()

    @staticmethod
    def _extract_identifier_name(document_info: dict[str, Any]) -> str:
        """
        Extract the core identifier name from document info.

        For modules: returns the module name
        For functions: returns the function name (without arity)

        Args:
            document_info: Document information dictionary

        Returns:
            The identifier name
        """
        if document_info["type"] == "module":
            return document_info["name"]
        else:
            return document_info["function"]

    def _initialize_bm25(self) -> tuple:
        """
        Initialize BM25 calculator with all documents in the index.

        Returns:
            Tuple of (BM25Okapi instance, document_map dict)
            - document_map maps document index to (type, module_name, location_info)
        """
        documents = []
        document_map = []

        # Collect all documents (modules and functions with keywords)
        for module_name, module_data in self.index.get("modules", {}).items():
            # Add module as a document
            if module_data.get("keywords"):
                doc_keywords = [kw.lower() for kw in module_data["keywords"]]
                documents.append(doc_keywords)
                document_map.append(
                    {
                        "type": "module",
                        "name": module_name,
                        "module": module_name,
                        "file": module_data["file"],
                        "line": module_data["line"],
                        "doc": module_data.get("moduledoc"),
                        "keywords": module_data["keywords"],
                    }
                )

            # Add functions as documents
            for func in module_data.get("functions", []):
                if func.get("keywords"):
                    doc_keywords = [kw.lower() for kw in func["keywords"]]
                    documents.append(doc_keywords)
                    full_name = f"{module_name}.{func['name']}/{func['arity']}"
                    document_map.append(
                        {
                            "type": "function",
                            "name": full_name,
                            "module": module_name,
                            "function": func["name"],
                            "arity": func["arity"],
                            "file": module_data["file"],
                            "line": func["line"],
                            "doc": func.get("doc"),
                            "keywords": func["keywords"],
                        }
                    )

        # If no documents were created (no keywords extracted), create documents using identifier names
        if not documents:
            for module_name, module_data in self.index.get("modules", {}).items():
                # Add module as a document using its name as keywords
                module_keywords = split_identifier(module_name)
                documents.append(module_keywords)
                document_map.append(
                    {
                        "type": "module",
                        "name": module_name,
                        "module": module_name,
                        "file": module_data["file"],
                        "line": module_data["line"],
                        "doc": module_data.get("moduledoc"),
                        "keywords": module_keywords,
                    }
                )

                # Add functions as documents
                for func in module_data.get("functions", []):
                    # Use extracted keywords if available, otherwise fall back to split identifier
                    if func.get("keywords"):
                        func_keywords = [kw.lower() for kw in func["keywords"]]
                    else:
                        func_keywords = split_identifier(func["name"])

                    documents.append(func_keywords)
                    full_name = f"{module_name}.{func['name']}/{func['arity']}"
                    document_map.append(
                        {
                            "type": "function",
                            "name": full_name,
                            "module": module_name,
                            "function": func["name"],
                            "arity": func["arity"],
                            "file": module_data["file"],
                            "line": func["line"],
                            "doc": func.get("doc"),
                            "keywords": func_keywords,
                        }
                    )

        # Initialize BM25 with all documents
        # Use b=0.4 (lower than default 0.75) to reduce length normalization penalty
        # This is appropriate for code search where longer names are more specific, not verbose
        bm25 = BM25Okapi(documents, b=0.4) if documents else None
        return bm25, document_map

    def _match_wildcard(self, pattern: str, text: str) -> bool:
        """
        Check if text matches a wildcard pattern.

        Supports * (matches any characters) only.

        Args:
            pattern: Wildcard pattern (e.g., "create*", "test_*")
            text: Text to match against

        Returns:
            True if text matches the pattern
        """
        # Only support * wildcard, not ?
        if "?" in pattern:
            return False
        return fnmatch.fnmatch(text.lower(), pattern.lower())

    def _expand_wildcard_keywords(
        self, query_keywords: list[str], document_keywords: list[str]
    ) -> list[str]:
        """
        Expand wildcard patterns to actual matching keywords from the document.

        Args:
            query_keywords: List of query keywords (may contain wildcards)
            document_keywords: List of keywords from a document

        Returns:
            List of actual matching keywords found in the document
        """
        matched_keywords = []
        for query_kw in query_keywords:
            for doc_kw in document_keywords:
                if self._match_wildcard(query_kw, doc_kw):
                    matched_keywords.append(query_kw)
        return matched_keywords

    def _expand_wildcard_keywords_with_identifier(
        self,
        query_keywords: list[str],
        document_keywords: list[str],
        identifier_name: str,
    ) -> list[str]:
        """
        Expand wildcard patterns to actual matching keywords from the document and identifier name.

        Args:
            query_keywords: List of query keywords (may contain wildcards)
            document_keywords: List of keywords from a document
            identifier_name: The full identifier name (function/module name)

        Returns:
            List of actual matching keywords found in the document or identifier
        """
        matched_keywords = []
        for query_kw in query_keywords:
            # Check against individual keywords
            for doc_kw in document_keywords:
                if self._match_wildcard(query_kw, doc_kw):
                    matched_keywords.append(query_kw)
                    break  # Only add each query keyword once

            # Also check against the full identifier name
            if query_kw not in matched_keywords and self._match_wildcard(query_kw, identifier_name):
                matched_keywords.append(query_kw)
        return matched_keywords

    def _get_wildcard_scores(self, query_keywords: list[str]) -> list[float]:
        """
        Calculate BM25-like scores for wildcard matching.

        Args:
            query_keywords: List of query keywords (may contain wildcards)

        Returns:
            List of scores for each document
        """
        scores = []

        for _, doc_info in enumerate(self.document_map):
            doc_keywords = [kw.lower() for kw in doc_info["keywords"]]
            identifier_name = self._extract_identifier_name(doc_info)

            # Find matching keywords using wildcard patterns
            # Check both individual keywords and full identifier name
            matched_keywords = self._expand_wildcard_keywords_with_identifier(
                query_keywords, doc_keywords, identifier_name
            )

            if matched_keywords:
                # Calculate a simple score based on number of matches
                # This is a simplified version of BM25 for wildcard matching
                score = len(matched_keywords) / len(query_keywords)
                scores.append(score)
            else:
                scores.append(0.0)

        return scores

    def _has_wildcards(self, keywords: list[str]) -> bool:
        """Check if any keywords contain wildcard patterns."""
        return any("*" in keyword for keyword in keywords)

    def search(self, query_keywords: list[str], top_n: int = 5) -> list[dict[str, Any]]:
        """
        Search for modules and functions matching the given keywords.

        Uses BM25 ranking to score documents based on keyword relevance.
        Identifier names (function/module names) are boosted significantly
        when they match query keywords.

        Automatically detects wildcard patterns (* supported) in keywords.

        Args:
            query_keywords: List of keywords to search for
            top_n: Maximum number of results to return

        Returns:
            List of result dictionaries sorted by score (descending), each containing:
            - type: 'module' or 'function'
            - name: Full name (e.g., 'MyApp.User' or 'MyApp.User.create/2')
            - module: Module name
            - file: File path
            - line: Line number
            - score: BM25 score (float), boosted if identifier matches query
            - confidence: Percentage of query keywords matched
            - matched_keywords: List of matched keywords
            - doc: Documentation string (if available)
        """
        if not query_keywords or self.bm25 is None or not self.document_map:
            return []

        # Normalize query keywords to lowercase
        query_keywords_lower = [kw.lower() for kw in query_keywords]

        # Check if wildcards are present
        enable_wildcards = self._has_wildcards(query_keywords_lower)

        # Get BM25 scores for all documents
        if enable_wildcards:
            # For wildcard matching, we need to manually score documents
            bm25_scores = self._get_wildcard_scores(query_keywords_lower)
        else:
            bm25_scores = self.bm25.get_scores(query_keywords_lower)

        results = []

        # Build results with scores
        for doc_idx, bm25_score in enumerate(bm25_scores):
            # BM25 can produce negative scores for small corpuses
            # We check matched keywords instead to filter relevance
            doc_info = self.document_map[doc_idx]

            # Check if there are any matching keywords first
            if enable_wildcards:
                identifier_name = self._extract_identifier_name(doc_info)
                matched = self._count_wildcard_matches(
                    query_keywords_lower, doc_info["keywords"], identifier_name
                )
            else:
                matched = self._count_matches(query_keywords_lower, doc_info["keywords"])

            # Only include documents that match at least one query keyword
            if matched["score"] > 0:
                # Apply identifier name boost
                if enable_wildcards:
                    final_score = self._apply_identifier_boost_wildcard(
                        bm25_score, query_keywords_lower, doc_info
                    )
                else:
                    final_score = self._apply_identifier_boost(
                        bm25_score, query_keywords_lower, doc_info
                    )

                # Apply name coverage penalty (penalize functions with extra words not in query)
                coverage_penalty = self._calculate_name_coverage_penalty(
                    query_keywords_lower, doc_info
                )
                # For negative scores, divide by penalty instead of multiply
                # This ensures penalty always makes the score worse regardless of sign
                if final_score < 0 and coverage_penalty < 1.0:
                    final_score = final_score / coverage_penalty
                else:
                    final_score = final_score * coverage_penalty

                result = {
                    "type": doc_info["type"],
                    "name": doc_info["name"],
                    "module": doc_info["module"],
                    "file": doc_info["file"],
                    "line": doc_info["line"],
                    "score": round(final_score, 4),
                    "confidence": matched["confidence"],
                    "matched_keywords": matched["matched_keywords"],
                }

                # Add type-specific fields
                if doc_info["type"] == "function":
                    result["function"] = doc_info["function"]
                    result["arity"] = doc_info["arity"]

                # Add documentation if available
                if doc_info.get("doc"):
                    result["doc"] = doc_info["doc"]

                results.append(result)

        # Sort by final score (descending), then by name for stable results
        results.sort(key=lambda x: (-x["score"], x["name"]))

        return results[:top_n]

    def _apply_identifier_boost(
        self, bm25_score: float, query_keywords: list[str], doc_info: dict[str, Any]
    ) -> float:
        """
        Apply boost to BM25 score if query keywords match the identifier name.

        Identifier names (function/module names) are given much higher weight.
        If any query keyword matches a word in the identifier name, the score
        is multiplied by IDENTIFIER_MATCH_BOOST.

        Args:
            bm25_score: The original BM25 score
            query_keywords: Normalized query keywords (lowercase)
            doc_info: Document information with function/module name

        Returns:
            Final score (boosted if identifier matches)
        """
        # Extract the identifier name (module name or function name)
        identifier_name = self._extract_identifier_name(doc_info)

        # Split identifier into words
        identifier_words = split_identifier(identifier_name)

        # Check if any query keyword matches any word in the identifier
        for query_kw in query_keywords:
            if query_kw in identifier_words:
                # Apply significant boost for identifier match
                return bm25_score * self.IDENTIFIER_MATCH_BOOST

        return bm25_score

    def _count_matches(self, query_keywords: list[str], item_keywords: list[str]) -> dict[str, Any]:
        """
        Count matching keywords between query and item.

        Args:
            query_keywords: Query keywords (normalized to lowercase)
            item_keywords: Keywords from module/function

        Returns:
            Dictionary with:
            - score: Number of matching keywords
            - confidence: Percentage match (score / len(query_keywords) * 100)
            - matched_keywords: List of matched keywords
        """
        # Normalize item keywords to lowercase
        item_keywords_lower = [kw.lower() for kw in item_keywords]

        # Find matches
        matched_keywords = []
        for query_kw in query_keywords:
            if query_kw in item_keywords_lower:
                matched_keywords.append(query_kw)

        score = len(matched_keywords)
        confidence = (score / len(query_keywords)) * 100 if query_keywords else 0

        return {
            "score": score,
            "confidence": round(confidence, 1),
            "matched_keywords": matched_keywords,
        }

    def _count_wildcard_matches(
        self,
        query_keywords: list[str],
        item_keywords: list[str],
        identifier_name: str | None = None,
    ) -> dict[str, Any]:
        """
        Count matching keywords between query and item using wildcard patterns.

        Args:
            query_keywords: Query keywords (may contain wildcards, normalized to lowercase)
            item_keywords: Keywords from module/function
            identifier_name: The full identifier name (function/module name)

        Returns:
            Dictionary with:
            - score: Number of matching keywords
            - confidence: Percentage match (score / len(query_keywords) * 100)
            - matched_keywords: List of matched keywords
        """
        # Normalize item keywords to lowercase
        item_keywords_lower = [kw.lower() for kw in item_keywords]

        # Find matches using wildcard patterns
        if identifier_name:
            matched_keywords = self._expand_wildcard_keywords_with_identifier(
                query_keywords, item_keywords_lower, identifier_name
            )
        else:
            matched_keywords = self._expand_wildcard_keywords(query_keywords, item_keywords_lower)

        score = len(matched_keywords)
        confidence = (score / len(query_keywords)) * 100 if query_keywords else 0

        return {
            "score": score,
            "confidence": round(confidence, 1),
            "matched_keywords": matched_keywords,
        }

    def _apply_identifier_boost_wildcard(
        self, bm25_score: float, query_keywords: list[str], doc_info: dict[str, Any]
    ) -> float:
        """
        Apply boost to BM25 score if query keywords match the identifier name using wildcards.

        Identifier names (function/module names) are given much higher weight.
        If any query keyword matches a word in the identifier name using wildcard patterns,
        the score is multiplied by IDENTIFIER_MATCH_BOOST.

        Args:
            bm25_score: The original BM25 score
            query_keywords: Query keywords (may contain wildcards, normalized to lowercase)
            doc_info: Document information with function/module name

        Returns:
            Final score (boosted if identifier matches)
        """
        # Extract the identifier name (module name or function name)
        identifier_name = self._extract_identifier_name(doc_info)

        # Split identifier into words
        identifier_words = split_identifier(identifier_name)

        # Check if any query keyword matches any word in the identifier using wildcards
        for query_kw in query_keywords:
            for identifier_word in identifier_words:
                if self._match_wildcard(query_kw, identifier_word):
                    # Apply significant boost for identifier match
                    return bm25_score * self.IDENTIFIER_MATCH_BOOST

        return bm25_score

    def _calculate_name_coverage_penalty(
        self, query_keywords: list[str], doc_info: dict[str, Any]
    ) -> float:
        """
        Calculate penalty for functions whose names contain words NOT in the query.

        This helps rank exact matches higher than functions with extra words in their names.
        For example, searching "create user" should rank "create_user" higher than
        "create_invalid_user" because "invalid" is not in the query.

        Args:
            query_keywords: Normalized query keywords (lowercase)
            doc_info: Document information with function/module name

        Returns:
            Penalty multiplier between 0.1 and 1.0:
            - 1.0 = no penalty (exact match or all extra words in query)
            - 0.7 = 1 extra word not in query (30% penalty)
            - 0.4 = 2 extra words not in query (60% penalty)
            - 0.1 = 3+ extra words not in query (90% penalty cap)
        """
        # Only apply to functions (not modules)
        if doc_info["type"] != "function":
            return 1.0

        # Get function name and split it
        func_name = doc_info["function"]
        func_words = set(split_identifier(func_name))

        # Find words in function name that are NOT in query
        query_set = set(query_keywords)
        extra_words = func_words - query_set

        # No penalty if all function name words are in query (exact match)
        if not extra_words:
            return 1.0

        # Apply 30% penalty per extra word, with a floor of 0.1 (max 90% penalty)
        # This penalty is strong enough to overcome BM25 length normalization bias
        penalty_per_word = 0.3
        total_penalty = min(len(extra_words) * penalty_per_word, 0.9)

        return 1.0 - total_penalty
