"""Tool search index using BM25 algorithm.

Provides lexical tool discovery for tiered tool description injection:
core tools are fully injected into the LLM prompt, while extended tools
are searchable on-demand via the ``tool_search`` tool.

Pure Python implementation — no external dependencies.
"""

from __future__ import annotations

import math
import re
from collections import Counter

from agentkit.tools.base import Tool

__all__ = ["ToolSearchIndex"]


_TOKEN_RE = re.compile(r"[a-zA-Z0-9_]+")


def _tokenize(text: str) -> list[str]:
    """Tokenize text into lowercase alphanumeric tokens.

    Splits on non-alphanumeric characters (keeping underscores so that
    snake_case tool names like ``read_file`` stay intact) and lowercases.
    Empty tokens are filtered out.
    """
    return [t for t in _TOKEN_RE.findall(text.lower()) if t]


class ToolSearchIndex:
    """BM25-based search index over :class:`Tool` descriptions.

    Each tool is indexed by its ``name`` + ``description`` + parameter
    metadata (from ``input_schema``) + ``tags``. Supports keyword search
    returning the most relevant tools ranked by BM25 score.

    Usage::

        index = ToolSearchIndex(tools)
        results = index.search("read file", top_k=5)
    """

    _DEFAULT_K1: float = 1.5
    _DEFAULT_B: float = 0.75

    def __init__(
        self,
        tools: list[Tool],
        k1: float = _DEFAULT_K1,
        b: float = _DEFAULT_B,
    ):
        if k1 < 0:
            raise ValueError(f"k1 must be >= 0, got {k1}")
        if not (0.0 <= b <= 1.0):
            raise ValueError(f"b must be in [0, 1], got {b}")

        self._tools: list[Tool] = list(tools)
        self._k1 = k1
        self._b = b

        # Build document corpus from tool metadata
        self._documents: list[str] = [self._tool_to_text(t) for t in self._tools]
        self._tokenized_docs: list[list[str]] = [_tokenize(doc) for doc in self._documents]
        self._doc_lengths: list[int] = [len(tokens) for tokens in self._tokenized_docs]
        self._avgdl: float = (
            sum(self._doc_lengths) / len(self._doc_lengths) if self._doc_lengths else 0.0
        )

        # Term frequencies per document
        self._term_freqs: list[Counter[str]] = [Counter(tokens) for tokens in self._tokenized_docs]

        # Document frequency per term (number of docs containing the term)
        self._doc_freq: Counter[str] = Counter()
        for tokens in self._tokenized_docs:
            for term in set(tokens):
                self._doc_freq[term] += 1

        # Precompute IDF for every term in the corpus
        self._idf: dict[str, float] = self._compute_idf()

    # ------------------------------------------------------------------
    # Index construction
    # ------------------------------------------------------------------

    @staticmethod
    def _tool_to_text(tool: Tool) -> str:
        """Convert a tool's searchable metadata into a single text document."""
        parts: list[str] = [str(tool.name), str(tool.description)]

        schema: dict[str, object] | None = tool.input_schema
        if schema:
            props = schema.get("properties", {})
            for pname, pinfo in props.items():
                parts.append(str(pname))
                if isinstance(pinfo, dict):
                    pdesc = pinfo.get("description", "")
                    if pdesc:
                        parts.append(str(pdesc))

        if tool.tags:
            parts.extend(str(t) for t in tool.tags)

        return " ".join(parts)

    def _compute_idf(self) -> dict[str, float]:
        """Compute IDF for each term using the BM25 formula.

        ``idf(t) = ln((N - n + 0.5) / (n + 0.5) + 1)``

        where ``N`` is the total number of documents and ``n`` is the
        number of documents containing term ``t``. The ``+1`` inside the
        logarithm guarantees a non-negative IDF.
        """
        n_docs = len(self._tools)
        if n_docs == 0:
            return {}

        idf: dict[str, float] = {}
        for term, df in self._doc_freq.items():
            idf[term] = math.log((n_docs - df + 0.5) / (df + 0.5) + 1.0)
        return idf

    # ------------------------------------------------------------------
    # Scoring
    # ------------------------------------------------------------------

    def _bm25_score(self, query_tokens: list[str], doc_idx: int) -> float:
        """Compute the BM25 score for a single document against the query."""
        if not query_tokens or doc_idx >= len(self._tools):
            return 0.0

        doc_len = self._doc_lengths[doc_idx]
        term_freq = self._term_freqs[doc_idx]
        norm = (
            (1 - self._b + self._b * (doc_len / self._avgdl)) if self._avgdl > 0 else (1 - self._b)
        )
        score = 0.0

        # Each unique query term contributes once
        for term in set(query_tokens):
            tf = term_freq.get(term, 0)
            if tf == 0:
                continue
            idf = self._idf.get(term, 0.0)
            if idf <= 0:
                continue
            denom = tf + self._k1 * norm
            score += idf * (tf * (self._k1 + 1)) / denom

        return score

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def search(self, query: str, top_k: int = 5) -> list[Tool]:
        """Search tools by keyword query.

        Args:
            query: Search keywords (natural language or tool-name fragments).
            top_k: Maximum number of tools to return.

        Returns:
            Tools ranked by relevance (most relevant first). Only tools
            with a positive BM25 score are returned.
        """
        if top_k <= 0 or not self._tools:
            return []

        query_tokens = _tokenize(query)
        if not query_tokens:
            return []

        scored: list[tuple[int, float]] = []
        for i in range(len(self._tools)):
            score = self._bm25_score(query_tokens, i)
            if score > 0:
                scored.append((i, score))

        scored.sort(key=lambda x: x[1], reverse=True)
        return [self._tools[i] for i, _ in scored[:top_k]]

    def __len__(self) -> int:
        return len(self._tools)
