"""Deferred-tool catalog for tool search: BM25 retrieval over deferrable tool
defs plus the budgeted, byte-stable catalog listing embedded in the bridge."""

from __future__ import annotations

import functools
import math
import re
import threading
from collections import Counter
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

import snowballstemmer

from tools.tool_search_names import TOOL_CALL_NAME, TOOL_DESCRIBE_NAME, TOOL_SEARCH_NAME

# Chars-per-token rule of thumb for budget estimates; 4.0 slightly
# underestimates, which is the safer direction (fewer false activations).
CHARS_PER_TOKEN = 4.0


@dataclass
class CatalogEntry:
    """One deferrable tool, in a form the bridge tools can search and serve."""

    name: str
    description: str
    schema: Dict[str, Any]  # The full {"type":"function", "function": {...}} entry.
    source: str  # "mcp" | "plugin" | "other"
    source_name: str  # Toolset name, e.g. "mcp-github" or "kanban"
    _tokens: List[str] = field(default_factory=list)  # pre-tokenized for BM25


_TOKEN_RE = re.compile(r"[A-Za-z0-9]+")

# Snowball stemmers carry mutable parsing state and bridge dispatch runs on
# parallel tool-call threads, so: one stemmer per thread, created lazily.
_thread_local = threading.local()


def _stemmer() -> Any:
    st = getattr(_thread_local, "stemmer", None)
    if st is None:
        st = snowballstemmer.stemmer("english")
        _thread_local.stemmer = st
    return st


@functools.lru_cache(maxsize=16384)
def _stem(token: str) -> str:
    """Stem one token, memoized across stateless catalog rebuilds."""
    return _stemmer().stemWord(token)


def _tokenize(text: str) -> List[str]:
    """Lowercase alphanumeric tokens, Snowball-stemmed (English).

    Shared by the index path and the query path so "issues" matches ``create_issue``.
    """
    if not text:
        return []
    return [_stem(token.lower()) for token in _TOKEN_RE.findall(text)]


def _fn(td: Dict[str, Any]) -> Dict[str, Any]:
    """The ``function`` block of a tool-def (``{}`` when absent/None)."""
    return td.get("function") or {}


def _registry_entry(name: str) -> Any:
    """Registry entry for ``name``; None when unregistered OR when the registry
    is unavailable/raises (lookup failures must never fail a bridge call).
    The import stays lazy: tests patch ``tools.registry.registry``."""
    try:
        from tools.registry import registry
        return registry.get_entry(name)
    except Exception:
        return None


def _entry_search_text(td: Dict[str, Any], source_label: str = "") -> str:
    """Search-text blob: split name words + source label + description +
    top-level parameter names. Schema bodies are excluded (noise, no recall
    gain). The ``mcp__`` prefix is dropped — it is in every MCP document, so
    its IDF is ~0. The source label lets a service-name query ("linear") reach
    a tool whose own name omits the vendor.
    """
    fn = _fn(td)
    name = fn.get("name", "")
    if name.startswith("mcp__"):
        name = name[len("mcp__"):]
    desc = fn.get("description", "") or ""
    params = ((fn.get("parameters") or {}).get("properties") or {})
    param_names = " ".join(params.keys())
    name_words = re.sub(r"[_.:-]", " ", name)
    extra = source_label if source_label and source_label not in name_words.split() else ""
    return f"{name_words} {extra} {desc} {param_names}"


def _classify_source(name: str) -> Tuple[str, str]:
    """Return (source_kind, source_name) for a registered tool name."""
    entry = _registry_entry(name)
    if entry is None:
        return ("other", "")
    try:
        return ("mcp" if entry.toolset.startswith("mcp-") else "plugin", entry.toolset)
    except Exception:  # malformed entry (no str toolset)
        return ("other", "")


def build_catalog(tool_defs: List[Dict[str, Any]]) -> List[CatalogEntry]:
    """Build the deferred-tool catalog from the deferrable subset of tool-defs."""
    catalog: List[CatalogEntry] = []
    for td in tool_defs:
        fn = _fn(td)
        name = fn.get("name", "")
        if not name:
            continue
        source, source_name = _classify_source(name)
        # Index the human-facing label ("linear", not "mcp-linear").
        source_label = _listing_group_label(source_name) if source_name else ""
        catalog.append(CatalogEntry(
            name=name,
            description=fn.get("description", "") or "",
            schema=td,
            source=source,
            source_name=source_name,
            _tokens=_tokenize(_entry_search_text(td, source_label)),
        ))
    return catalog


def _bm25_score(query_tokens: List[str], doc_tokens: List[str],
                doc_lengths: List[int], avg_dl: float,
                doc_freq: Dict[str, int], n_docs: int,
                k1: float = 1.5, b: float = 0.75) -> float:
    """Standard BM25 score for one query against one document (inlined; the
    catalog is bounded — typically < 500 tools — so a dependency is not worth it)."""
    if not doc_tokens:
        return 0.0
    score = 0.0
    dl = len(doc_tokens)
    doc_tf = Counter(doc_tokens)
    for q in query_tokens:
        df = doc_freq.get(q, 0)
        if df == 0:
            continue
        idf = math.log(1 + (n_docs - df + 0.5) / (df + 0.5))
        tf = doc_tf.get(q, 0)
        if tf == 0:
            continue
        norm = tf * (k1 + 1) / (tf + k1 * (1 - b + b * dl / max(avg_dl, 1.0)))
        score += idf * norm
    return score


_CorpusStats = Tuple[List[int], float, Dict[str, int], int]


def _corpus_stats(catalog: List[CatalogEntry]) -> _CorpusStats:
    """Compute the BM25 statistics shared by every query over a catalog."""
    doc_lengths = [len(entry._tokens) for entry in catalog]
    avg_dl = sum(doc_lengths) / max(len(doc_lengths), 1)
    doc_freq: Dict[str, int] = Counter()
    for entry in catalog:
        doc_freq.update(set(entry._tokens))
    return doc_lengths, avg_dl, dict(doc_freq), len(catalog)


def search_catalog(
    catalog: List[CatalogEntry],
    query: str,
    limit: int = 5,
    *,
    corpus_stats: Optional[_CorpusStats] = None,
) -> List[CatalogEntry]:
    """Top-``limit`` catalog entries for ``query`` by BM25 (exact name match
    ranks first). Falls back to a name-substring match only when NO query
    token appears in any document (e.g. "hub" vs ``github_*``); the IDF
    variant is strictly positive, so a hit anywhere suppresses the fallback.
    """
    if not catalog or limit <= 0:
        return []
    query_tokens = _tokenize(query)
    if not query_tokens:
        return []

    if corpus_stats is None:
        corpus_stats = _corpus_stats(catalog)
    doc_lengths, avg_dl, doc_freq, n_docs = corpus_stats

    scored: List[Tuple[float, CatalogEntry]] = []
    exact_name = query.strip().lower()
    for entry in catalog:
        if entry.name.lower() == exact_name:
            scored.append((float("inf"), entry))
            continue
        s = _bm25_score(query_tokens, entry._tokens, doc_lengths, avg_dl,
                        doc_freq, n_docs)
        if s > 0:
            scored.append((s, entry))

    if not scored:
        ql = query.lower()
        for entry in catalog:
            if ql in entry.name.lower():
                scored.append((0.1, entry))

    scored.sort(key=lambda x: x[0], reverse=True)
    return [e for _, e in scored[:limit]]


# Sentence end: ., !, ? followed by whitespace/EOS, not inside e.g./i.e./etc.
_SENTENCE_END_RE = re.compile(r"(?<!\be\.g)(?<!\bi\.e)(?<!\betc)[.!?](?=\s|$)")


def _short_desc(description: str, max_chars: int = 60) -> str:
    """First sentence of a tool description, clipped to ``max_chars`` on a
    word boundary. ``e.g.``/``i.e.``/``etc.`` do not end a sentence; whitespace
    normalization and the regex search stay linear-time on hostile input."""
    text = " ".join((description or "").split())
    if not text:
        return ""
    m = _SENTENCE_END_RE.search(text)
    if m:
        text = text[:m.end()]
    if len(text) <= max_chars:
        return text
    clipped = text[:max_chars]
    if " " in clipped:
        clipped = clipped.rsplit(" ", 1)[0]
    return clipped.rstrip(",;: ") + "…"


def _listing_group_label(source_name: str) -> str:
    """Human-facing group heading for a toolset, e.g. ``mcp-github`` -> ``github``."""
    label = source_name or "other"
    return label[4:] if label.startswith("mcp-") else label


def build_catalog_listing(
    deferrable: List[Dict[str, Any]],
    *,
    max_tokens: int = 4000,
) -> Optional[str]:
    """Render the deferred-catalog manifest; text only (see
    :func:`build_catalog_listing_with_form` for the degradation ladder)."""
    return build_catalog_listing_with_form(deferrable, max_tokens=max_tokens)[0]


def build_catalog_listing_with_form(
    deferrable: List[Dict[str, Any]],
    *,
    max_tokens: int = 4000,
) -> Tuple[Optional[str], str]:
    """Render the skills-style deferred-catalog manifest: ``- name: short desc``
    lines grouped under a heading per source (MCP server / plugin toolset).

    Returns ``(text, form)``; ``form`` is ``"full"``, ``"names"`` (names-only),
    ``"mixed"`` (oversized servers collapsed to a name + count summary line,
    small ones keep per-tool lines), ``"groups"`` (every server summarized),
    or ``"none"`` (over budget even summarized -> text is None).

    Ordering is deterministic (sorted groups and tools) so the block is
    byte-stable across assemblies — the request prefix stays cacheable.
    Degradation is PER SERVER (largest first): one huge server must not cost
    a small co-attached server its listing.
    """
    if not deferrable:
        return None, "none"

    groups: Dict[str, List[Tuple[str, str]]] = {}
    for td in deferrable:
        fn = _fn(td)
        name = fn.get("name", "")
        if not name:
            continue
        # ``_classify_source`` returns ("other", "") for unregistered names and
        # ``_listing_group_label("")`` is "other", so one call covers both.
        label = _listing_group_label(_classify_source(name)[1])
        groups.setdefault(label, []).append((name, _short_desc(fn.get("description", ""))))

    if not groups:
        return None, "none"

    def render_group(label: str, mode: str) -> str:
        """Render one server's block. mode: 'full' | 'names' | 'summary'."""
        tools = sorted(groups[label])
        if mode == "summary":
            return (f"{label} ({len(tools)} tools — names not listed; "
                    f"discover via `{TOOL_SEARCH_NAME}`)")
        lines = [f"{label} tools ({len(tools)}):"]
        if mode == "full":
            lines.extend(f"- {name}: {desc}" if desc else f"- {name}" for name, desc in tools)
        else:
            lines.append(", ".join(name for name, _ in tools))
        return "\n".join(lines)

    header = ("Deferred tool catalog (call schemas via "
              f"`{TOOL_DESCRIBE_NAME}`, invoke via `{TOOL_CALL_NAME}`):")

    def assemble_if_fits(modes: Dict[str, str]) -> Optional[str]:
        text = "\n".join([header] + [render_group(lbl, modes[lbl]) for lbl in sorted(groups)])
        return text if math.ceil(len(text) / CHARS_PER_TOKEN) <= max_tokens else None

    # 1. Everything full.  2. Everything names-only.
    for mode in ("full", "names"):
        modes = {lbl: mode for lbl in groups}
        text = assemble_if_fits(modes)
        if text is not None:
            return text, mode

    # 3. Per-server degradation: collapse the LARGEST rendered groups first
    #    (deterministic: size then label) so one oversized server does not
    #    cost a small co-attached server its per-tool names.
    by_size = sorted(groups, key=lambda lbl: (-len(render_group(lbl, "names")), lbl))
    for lbl in by_size:
        modes[lbl] = "summary"
        text = assemble_if_fits(modes)
        if text is not None:
            form = "groups" if all(m == "summary" for m in modes.values()) else "mixed"
            return text, form

    return None, "none"