# Copyright (C) 2026 GoodData Corporation

"""Hybrid scoring (lexical + embeddings) for rules retrieval."""

from __future__ import annotations

import logging

from aida_mcp.tools.rules.config import (
    EMBEDDING_PENALTY,
    MIN_EMBEDDING_FOR_GENERIC_PAIRS,
    MIN_EMBEDDING_FOR_SINGLE_MATCH,
    MIN_EMBEDDING_SCORE,
    PHRASE_MATCH_BOOST,
)
from aida_mcp.tools.rules.embeddings.chunker import RuleChunk
from aida_mcp.tools.rules.embeddings.store import VectorStore
from aida_mcp.tools.rules.scoring_bm25 import BM25ChunkScorer, DescriptionBooster
from aida_mcp.tools.rules.scoring_embedding import EmbeddingScorer
from aida_mcp.tools.rules.scoring_types import ScoredRule, tokenize_query
from aida_mcp.tools.rules.store import Rule

logger = logging.getLogger(__name__)


class HybridScorer:
    """Combine lexical + embedding rankings using Reciprocal Rank Fusion."""

    RRF_K_MIN: int = 5
    RRF_K_MAX: int = 60

    HIGH_CHUNK_SCORE_THRESHOLD: float = 0.55
    HIGH_CHUNK_SCORE_BOOST: float = 0.15

    def __init__(self, rules: list[Rule], chunks: list[RuleChunk], vector_store: VectorStore | None = None):
        self._rules = rules
        self._chunks = chunks
        self._bm25_chunks = BM25ChunkScorer(chunks)
        self._booster = DescriptionBooster()
        self._embedding_scorer = EmbeddingScorer(vector_store) if vector_store else None
        self._rrf_k = self._compute_rrf_k(len(rules))

    def _compute_rrf_k(self, num_items: int) -> int:
        dynamic_k = num_items // 5
        return max(self.RRF_K_MIN, min(self.RRF_K_MAX, dynamic_k))

    def score(self, query: str) -> list[ScoredRule]:
        query_tokens = tokenize_query(query)
        query_term_count = len(query_tokens)

        rule_by_path: dict[str, Rule] = {r.path: r for r in self._rules}

        bm25_data: dict[str, tuple[float, set[str], bool]] = {}
        bm25_chunks_by_path: dict[str, list[tuple[float, RuleChunk]]] = {}

        for chunk in self._chunks:
            score, matched, has_phrase = self._bm25_chunks.score(query_tokens, chunk)
            if score <= 0:
                continue

            bm25_chunks_by_path.setdefault(chunk.rule_path, []).append((score, chunk))
            prev = bm25_data.get(chunk.rule_path)
            if prev is None or score > prev[0]:
                bm25_data[chunk.rule_path] = (score, matched, has_phrase)

        for path, (score, matched, has_phrase) in list(bm25_data.items()):
            rule = rule_by_path.get(path)
            if rule is None:
                continue
            boosted_score = self._booster.boost(query, rule, score)
            bm25_data[path] = (boosted_score, matched, has_phrase)

        embedding_scores: dict[str, float] = {}
        matched_chunks: dict[str, list[RuleChunk]] = {}
        if self._embedding_scorer:
            embedding_scores, matched_chunks = self._embedding_scorer.score_rules(query, self._rules)

        bm25_ranking = self._scores_to_ranks({p: s for p, (s, _, _) in bm25_data.items()})
        emb_ranking = self._scores_to_ranks(embedding_scores) if embedding_scores else {}

        results: list[ScoredRule] = []
        for rule in self._rules:
            bm25_score, matched, has_phrase = bm25_data.get(rule.path, (0.0, set(), False))
            emb_score = embedding_scores.get(rule.path, 0.0)
            chunks = matched_chunks.get(rule.path, [])

            if not chunks:
                scored_chunks = sorted(bm25_chunks_by_path.get(rule.path, []), key=lambda x: x[0], reverse=True)
                chunks = [c for _, c in scored_chunks[:3]]

            # Apply minimum match requirements for multi-term queries
            # Only apply strict filtering when embeddings are available (to avoid breaking BM25-only tests)
            if query_term_count >= 2 and self._embedding_scorer is not None:
                matched_count = len(matched)
                # Require at least 2 matches OR 1 strong phrase match OR high embedding score
                # Single word match on multi-term query - skip if no phrase and embedding below stricter threshold
                if matched_count == 1 and not has_phrase and emb_score < MIN_EMBEDDING_FOR_SINGLE_MATCH:
                    # Skip this rule - insufficient matches (likely false positive)
                    continue
                # For 2-word matches without phrase - require higher embedding score to filter false positives
                if matched_count == 2 and not has_phrase and emb_score < MIN_EMBEDDING_FOR_GENERIC_PAIRS:
                    # Two-word match without phrase context and low embedding - likely false positive
                    continue

            rrf_score = 0.0
            bm25_rank = bm25_ranking.get(rule.path, len(self._rules))
            rrf_score += 1.0 / (self._rrf_k + bm25_rank)

            if emb_ranking:
                emb_rank = emb_ranking.get(rule.path, len(self._rules))
                rrf_score += 1.0 / (self._rrf_k + emb_rank)

            # Phrase matching boost
            if has_phrase:
                rrf_score += PHRASE_MATCH_BOOST
                logger.debug("Phrase match boost +%.2f for %s", PHRASE_MATCH_BOOST, rule.path)

            if emb_score >= self.HIGH_CHUNK_SCORE_THRESHOLD:
                rrf_score += self.HIGH_CHUNK_SCORE_BOOST
                logger.debug(
                    "Boosting %s by %.2f (chunk score %.2f >= %.2f)",
                    rule.path,
                    self.HIGH_CHUNK_SCORE_BOOST,
                    emb_score,
                    self.HIGH_CHUNK_SCORE_THRESHOLD,
                )

            # Embedding score floor penalty for multi-term queries
            if query_term_count >= 2 and emb_score < MIN_EMBEDDING_SCORE and rrf_score > 0:
                rrf_score *= 1.0 - EMBEDDING_PENALTY
                logger.debug(
                    "Embedding penalty -%.0f%% for %s (score %.2f < %.2f)",
                    EMBEDDING_PENALTY * 100,
                    rule.path,
                    emb_score,
                    MIN_EMBEDDING_SCORE,
                )

            results.append(
                ScoredRule(
                    rule=rule,
                    score=rrf_score,
                    matched_tokens=matched,
                    bm25_score=bm25_score,
                    embedding_score=emb_score,
                    matched_chunks=chunks,
                )
            )

        results.sort(key=lambda x: x.score, reverse=True)
        return results

    def _scores_to_ranks(self, scores: dict[str, float]) -> dict[str, int]:
        sorted_items = sorted(scores.items(), key=lambda x: x[1], reverse=True)
        ranks: dict[str, int] = {}
        current_rank = 1
        prev_score: float | None = None

        for i, (path, score) in enumerate(sorted_items):
            if prev_score is not None and score < prev_score:
                current_rank = i + 1
            ranks[path] = current_rank
            prev_score = score

        return ranks
