# Copyright (C) 2026 GoodData Corporation

"""Lexical (BM25-like) scoring for rules retrieval."""

from __future__ import annotations

import math
import re
from dataclasses import dataclass

from aida_mcp.tools.rules.config import PHRASE_PROXIMITY
from aida_mcp.tools.rules.embeddings.chunker import RuleChunk
from aida_mcp.tools.rules.store import Rule


class BM25ChunkScorer:
    """BM25 scoring for section-level chunks.

    Uses chunk.tokens and chunk.boost_tokens to simulate higher TF for anchor terms
    (headings, identifiers, paths) without introducing a manual alias list.
    """

    def __init__(self, chunks: list[RuleChunk], k1: float = 1.5, b: float = 0.75):
        self._chunks = chunks
        self._k1 = k1
        self._b = b

        self._all_tokens: set[str] = set()
        for c in chunks:
            self._all_tokens.update(c.tokens)

        self._idf_scores = self._compute_idf_scores()
        self._avg_doc_length = sum(len(c.tokens) for c in chunks) / len(chunks) if chunks else 1

    def _compute_idf_scores(self) -> dict[str, float]:
        n_docs = len(self._chunks)
        idf_scores: dict[str, float] = {}
        for token in self._all_tokens:
            doc_freq = sum(1 for c in self._chunks if token in c.tokens)
            if doc_freq > 0:
                idf_scores[token] = math.log((n_docs - doc_freq + 0.5) / (doc_freq + 0.5) + 1)
        return idf_scores

    def score(self, query_tokens: set[str], chunk: RuleChunk) -> tuple[float, set[str], bool]:
        """Score a chunk, returning (score, matched_tokens, has_phrase_match)."""
        score = 0.0
        matched: set[str] = set()
        doc_tokens = chunk.tokens
        doc_length = len(doc_tokens)

        for token in query_tokens:
            if token not in doc_tokens:
                continue

            matched.add(token)

            # Binary TF baseline, slightly boosted for anchor tokens (headings/identifiers/paths).
            tf = 2 if token in chunk.boost_tokens else 1
            idf = self._idf_scores.get(token, 1.0)

            length_norm = 1 - self._b + self._b * (doc_length / self._avg_doc_length) if self._avg_doc_length > 0 else 1
            numerator = tf * (self._k1 + 1)
            denominator = tf + self._k1 * length_norm
            score += idf * (numerator / denominator)

        # Slightly down-weight description-only chunks so headings/sections win when available.
        if chunk.is_description:
            score *= 0.7

        # Check for phrase match: query tokens appearing close together in chunk text
        has_phrase = self._has_phrase_match(query_tokens, chunk.text)

        return score, matched, has_phrase

    def _has_phrase_match(self, query_tokens: set[str], text: str) -> bool:
        """Check if query tokens appear together (within PHRASE_PROXIMITY words) in text."""
        if len(query_tokens) < 2:
            return False

        # Tokenize text (simple word-based)
        words = text.lower().split()
        if len(words) < 2:
            return False

        # Find positions of query tokens
        token_positions: dict[str, list[int]] = {}
        for i, word in enumerate(words):
            # Simple word matching (handle punctuation)
            clean_word = "".join(c for c in word if c.isalnum())
            if clean_word in query_tokens:
                token_positions.setdefault(clean_word, []).append(i)

        # Check if any two query tokens appear within proximity
        token_list = list(query_tokens)
        for i in range(len(token_list)):
            for j in range(i + 1, len(token_list)):
                token1, token2 = token_list[i], token_list[j]
                pos1 = token_positions.get(token1, [])
                pos2 = token_positions.get(token2, [])

                for p1 in pos1:
                    for p2 in pos2:
                        if abs(p1 - p2) <= PHRASE_PROXIMITY:
                            return True

        return False


@dataclass(frozen=True)
class DescriptionBooster:
    """Boosts scores based on description matches."""

    SKIP_WORDS: frozenset[str] = frozenset(
        {
            "the",
            "for",
            "and",
            "how",
            "what",
            "why",
            "when",
            "where",
            "does",
            "service",
            "microservice",
            "about",
            "with",
            "from",
        }
    )

    def boost(self, query: str, rule: Rule, base_score: float) -> float:
        query_lower = query.lower()
        desc_lower = rule.description.lower()

        boost = 1.0

        # Extract meaningful terms from query
        query_terms = set(re.findall(r"[a-z][a-z0-9-]*[a-z0-9]", query_lower))
        query_terms -= set(self.SKIP_WORDS)

        for term in query_terms:
            if term in desc_lower:
                if "-" in term:
                    boost += 10.0
                elif len(term) > 5:
                    boost += 2.0
                else:
                    boost += 0.5

        return base_score * boost
