"""
Semantic Search — enhanced module scoring beyond exact keyword matching.

Problem with pure keyword scoring:
  Query "where is payment handled?" misses modules named "billing", "checkout",
  "transaction" — because none of those words appear in the query.

Solution — two-tier approach:
  Tier 1 (always, zero deps):
    - Light stemming (strips -ing, -tion, -ers, etc.) via regex
    - IDF weighting (rare domain terms score higher than common words)
    - Bigram matching (e.g. "context_engine" matches "context engine")
    - Camelcase/snake_case splitting ("getUserById" → ["get","user","by","id"])

  Tier 2 (optional, requires numpy):
    - TF-IDF vector cosine similarity computed at query time
    - 60% keyword score + 40% semantic similarity blending
    - No external ML model needed — pure numpy math

This replaces the simple keyword loop in ContextEngine._score_modules().
"""

from __future__ import annotations

import re
import math
from typing import Optional


# ── Tokenizer ─────────────────────────────────────────────────────────────────

_CAMEL_RE = re.compile(r"(?<=[a-z0-9])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])")


def tokenize(text: str) -> list[str]:
    """
    Split text into lowercase tokens, handling:
      - snake_case, kebab-case, dot.notation
      - camelCase and PascalCase identifiers
      - regular spaces and punctuation
    """
    # Split camelCase first
    text = _CAMEL_RE.sub(" ", text)
    # Then split on any non-alphanumeric char
    tokens = re.split(r"[^a-z0-9]+", text.lower())
    return [t for t in tokens if len(t) >= 2]


# ── Stemmer ───────────────────────────────────────────────────────────────────

_SUFFIX_RULES: list[tuple[str, str]] = [
    (r"ations?$", "ate"),
    (r"tions?$", "te"),
    (r"ings?$", ""),
    (r"ers?$", ""),
    (r"ies$", "y"),
    (r"ness$", ""),
    (r"ments?$", ""),
    (r"ists?$", ""),
    (r"s$", ""),
]


def stem(word: str) -> str:
    """Lightweight rule-based stemmer — no external dependencies."""
    if len(word) <= 4:
        return word
    for pattern, replacement in _SUFFIX_RULES:
        if re.search(pattern, word):
            return re.sub(pattern, replacement, word)
    return word


def expand_tokens(tokens: list[str]) -> set[str]:
    """Return original tokens + their stems."""
    expanded = set(tokens)
    for t in tokens:
        expanded.add(stem(t))
    return expanded


# ── Module text extraction ─────────────────────────────────────────────────────

def module_tokens(mod: dict) -> list[str]:
    """Extract all meaningful tokens from a module record."""
    parts = [
        mod.get("module_name", ""),
        mod.get("summary", ""),
        " ".join(mod.get("classes", [])),
        " ".join(mod.get("functions", [])),
        " ".join(mod.get("exports", [])),
        mod.get("path", "").replace("\\", "/"),
    ]
    return tokenize(" ".join(p for p in parts if p))


# ── IDF computation ───────────────────────────────────────────────────────────

def compute_idf(modules: dict) -> dict[str, float]:
    """
    Compute Inverse Document Frequency across all modules.

    Rare tokens (appear in few modules) get high IDF → higher relevance boost.
    Common tokens (appear everywhere) get IDF ≈ 1 → minimal boost.
    """
    n = len(modules)
    if n == 0:
        return {}

    df: dict[str, int] = {}
    for mod in modules.values():
        seen = set(module_tokens(mod))
        for token in seen:
            df[token] = df.get(token, 0) + 1
        # Also index stems
        for token in list(seen):
            s = stem(token)
            if s != token:
                df[s] = df.get(s, 0) + 1

    return {
        token: math.log((n + 1) / (count + 1)) + 1.0
        for token, count in df.items()
    }


# ── TF-IDF vector helpers (used for cosine similarity, numpy optional) ────────

def build_vocab(idf: dict[str, float]) -> list[str]:
    """Return sorted vocabulary list (for consistent vector indices)."""
    return sorted(idf.keys())


def vectorize(tokens: list[str], vocab_idx: dict[str, int], idf: dict[str, float]) -> list[float]:
    """Build a TF-IDF vector for a token list."""
    tf: dict[str, int] = {}
    for t in tokens:
        tf[t] = tf.get(t, 0) + 1
    n = max(len(tokens), 1)

    vec = [0.0] * len(vocab_idx)
    for token, count in tf.items():
        expanded = {token, stem(token)}
        for t in expanded:
            if t in vocab_idx:
                vec[vocab_idx[t]] = max(vec[vocab_idx[t]], (count / n) * idf.get(t, 1.0))
    return vec


# ── Main scorer ───────────────────────────────────────────────────────────────

class SemanticScorer:
    """
    Drop-in replacement for ContextEngine's keyword scorer.

    Tier 1 (always active): stemmed keyword scoring + IDF weighting
    Tier 2 (when numpy available): TF-IDF cosine similarity blended in
    """

    def __init__(self, modules: dict):
        self.modules = modules
        self.idf = compute_idf(modules)
        self._has_numpy = self._check_numpy()

        # Pre-compute per-module token sets for fast scoring
        self._module_token_sets: dict[str, set[str]] = {}
        self._module_tokens_raw: dict[str, list[str]] = {}
        for path, mod in modules.items():
            tokens = module_tokens(mod)
            self._module_tokens_raw[path] = tokens
            self._module_token_sets[path] = expand_tokens(tokens)

    def _check_numpy(self) -> bool:
        try:
            import numpy  # noqa: F401
            return True
        except ImportError:
            return False

    def score(self, query: str, top_k: Optional[int] = None) -> list[tuple[float, str]]:
        """
        Score all modules against the query.

        Returns: sorted list of (score, path), highest score first.
        """
        query_tokens = tokenize(query)
        query_expanded = expand_tokens(query_tokens)
        query_lower = query.lower()

        scored: list[tuple[float, str]] = []

        for path, mod in self.modules.items():
            score = self._score_tier1(mod, path, query_expanded, query_lower)

            if score > 0:
                scored.append((score, path))

        # Tier 2: blend cosine similarity if numpy available
        if self._has_numpy and scored:
            scored = self._blend_cosine(query_tokens, scored)

        scored.sort(key=lambda x: -x[0])
        return scored[:top_k] if top_k else scored

    def _score_tier1(
        self,
        mod: dict,
        path: str,
        query_expanded: set[str],
        query_phrase: str,
    ) -> float:
        """Tier 1: IDF-weighted stemmed keyword scoring with prefix matching."""
        score = 0.0
        module_name = mod.get("module_name", "").lower()
        summary = mod.get("summary", "").lower()
        mod_token_set = self._module_token_sets.get(path, set())

        for qtoken in query_expanded:
            if len(qtoken) < 2:
                continue
            idf_w = self.idf.get(qtoken, 1.0)

            # Exact + prefix/contains matching against module tokens
            # e.g. "authentication" matches "auth", "authenticate" matches "auth"
            def _matches(candidate_tokens: list[str]) -> bool:
                for ct in candidate_tokens:
                    if ct == qtoken:
                        return True
                    # prefix: "auth" matches "authentication" and vice versa
                    if len(qtoken) >= 4 and len(ct) >= 4:
                        if ct.startswith(qtoken) or qtoken.startswith(ct):
                            return True
                return False

            if _matches(tokenize(module_name)):
                score += 10 * idf_w
            for cls in mod.get("classes", []):
                if _matches(tokenize(cls)):
                    score += 8 * idf_w
            for fn in mod.get("functions", []):
                if _matches(tokenize(fn)):
                    score += 5 * idf_w
            if _matches(tokenize(path)):
                score += 4 * idf_w
            if _matches(tokenize(summary)):
                score += 3 * idf_w
            for imp in mod.get("imports", []):
                if _matches(tokenize(imp)):
                    score += 1 * idf_w

        # Phrase bonuses
        if query_phrase in summary:
            score += 5.0
        if query_phrase in module_name:
            score += 5.0

        return score

    def _blend_cosine(
        self,
        query_tokens: list[str],
        keyword_scored: list[tuple[float, str]],
    ) -> list[tuple[float, str]]:
        """
        Tier 2: blend keyword scores with TF-IDF cosine similarity.
        Ratio: 60% keyword + 40% cosine — keeps keyword signal dominant.
        """
        try:
            import numpy as np

            vocab = build_vocab(self.idf)
            vocab_idx = {w: i for i, w in enumerate(vocab)}

            # Query vector
            q_vec = np.array(vectorize(query_tokens, vocab_idx, self.idf), dtype=float)
            q_norm = np.linalg.norm(q_vec)
            if q_norm == 0:
                return keyword_scored
            q_unit = q_vec / q_norm

            max_kw = max(s for s, _ in keyword_scored) or 1.0
            blended = []

            for kw_score, path in keyword_scored:
                mod_vec = np.array(
                    vectorize(self._module_tokens_raw.get(path, []), vocab_idx, self.idf),
                    dtype=float,
                )
                mod_norm = np.linalg.norm(mod_vec)
                if mod_norm > 0:
                    cos_sim = float(np.dot(q_unit, mod_vec / mod_norm))
                else:
                    cos_sim = 0.0

                # Normalise keyword score to [0, max_kw] then blend
                blended_score = 0.6 * kw_score + 0.4 * cos_sim * max_kw
                blended.append((blended_score, path))

            return blended
        except Exception:
            return keyword_scored
