# Copyright (C) 2026 GoodData Corporation

"""Rules retriever for querying and scoring development rules."""

from __future__ import annotations

from pathlib import Path

from aida_mcp.tools.rules.config import (
    DEFAULT_SCORE_THRESHOLD,
    DEFAULT_TOP_K,
    EXPANDED_TOP_K,
    EXPANSION_TRIGGER_COUNT,
    FALLBACK_MIN_EMBEDDING,
    FALLBACK_MIN_SCORE,
    HIGH_SCORE_THRESHOLD,
    MAX_TOP_K,
    MIN_TOP_K,
    SCORE_GAP_THRESHOLD,
)
from aida_mcp.tools.rules.embeddings.chunker import RuleChunker
from aida_mcp.tools.rules.embeddings.store import VectorStore
from aida_mcp.tools.rules.scoring import HybridScorer, ScoredRule, tokenize_query
from aida_mcp.tools.rules.store import RuleStore, get_rule_store


class RulesRetriever:
    """
    Retrieves relevant development rules based on user queries.

    Uses hybrid scoring (BM25 + embeddings) when embeddings are available,
    falls back to BM25-only when not initialized.
    """

    def __init__(
        self,
        rules_dir: Path,
        *,
        score_threshold: float = DEFAULT_SCORE_THRESHOLD,
        use_embeddings: bool = True,
        vector_store: VectorStore | None = None,
    ):
        """
        Initialize the rules retriever.

        Args:
            rules_dir: Path to the rules directory
            score_threshold: Minimum score to include a rule
            use_embeddings: Whether to use embedding-based scoring (if available)
            vector_store: Optional vector store already scoped to this rules directory
        """
        self.rules_dir = rules_dir
        self.score_threshold = score_threshold
        self.use_embeddings = use_embeddings
        self._vector_store = vector_store

        # Initialize components
        self._store = get_rule_store(rules_dir)
        self._hybrid_scorer: HybridScorer | None = None

    @property
    def store(self) -> RuleStore:
        """Get the rule store."""
        return self._store

    def _get_scorer(self) -> HybridScorer:
        """Get or create the hybrid scorer."""
        if self._hybrid_scorer is None:
            # Chunk rules into markdown sections for both lexical and embedding matching.
            chunker = RuleChunker()
            chunks = chunker.chunk_rules(self._store.get_all_rules())
            self._hybrid_scorer = HybridScorer(
                self._store.get_all_rules(),
                chunks,
                vector_store=self._vector_store if self.use_embeddings else None,
            )
        return self._hybrid_scorer

    def retrieve(
        self,
        query: str,
        *,
        include_always_apply: bool = True,
        expand_required: bool = True,
        inject_paths: list[str] | None = None,
    ) -> list[ScoredRule]:
        """
        Retrieve relevant rules for a query.

        Args:
            query: The search query (user's question or task description)
            include_always_apply: Whether to always include alwaysApply rules
            expand_required: Whether to auto-include required dependencies
            inject_paths: Optional list of rule paths to inject (from bundles)

        Returns:
            List of scored rules, sorted by relevance
        """
        all_rules = self._store.get_all_rules()
        if not all_rules:
            return []

        # Check for empty query
        query_tokens = tokenize_query(query)
        if not query_tokens:
            if include_always_apply:
                return [ScoredRule(r, 1.0, set()) for r in self._store.get_always_apply_rules()]
            return []

        # Score all rules using hybrid scorer
        scorer = self._get_scorer()
        scored_rules = scorer.score(query)
        scored_by_path = {sr.rule.path: sr for sr in scored_rules}

        # Collect results
        results: list[ScoredRule] = []
        seen_paths: set[str] = set()

        # First, add always-apply rules if requested
        if include_always_apply:
            for sr in scored_rules:
                if sr.rule.always_apply and sr.rule.path not in seen_paths:
                    sr.selection_reasons = ["always"]
                    results.append(sr)
                    seen_paths.add(sr.rule.path)

        # Determine adaptive top_k based on score distribution
        # Count high-scoring rules to decide if we should expand
        high_score_count = sum(
            1
            for sr in scored_rules
            if sr.rule.path not in seen_paths and sr.score >= HIGH_SCORE_THRESHOLD and sr.score >= self.score_threshold
        )

        # Use expanded top_k if many high-scoring rules (complex query)
        adaptive_top_k = EXPANDED_TOP_K if high_score_count >= EXPANSION_TRIGGER_COUNT else DEFAULT_TOP_K
        adaptive_top_k = min(adaptive_top_k, MAX_TOP_K)  # Cap at maximum

        # Then add top-k scored rules above threshold (high-confidence)
        # Use score gap detection to stop early if relevance drops significantly
        previous_score: float | None = None
        for sr in scored_rules:
            if len(results) >= adaptive_top_k:
                break
            if sr.rule.path not in seen_paths and sr.score >= self.score_threshold:
                # Score gap detection: stop if score drops significantly from previous
                if previous_score is not None and len(results) >= MIN_TOP_K:
                    score_drop = (previous_score - sr.score) / previous_score
                    if score_drop > SCORE_GAP_THRESHOLD:
                        # Significant drop - stop here (natural cutoff)
                        break

                if "score" not in sr.selection_reasons:
                    sr.selection_reasons.append("score")
                results.append(sr)
                seen_paths.add(sr.rule.path)
                previous_score = sr.score

        # Best-effort fallback: if we still don't have enough results, fill with next best rules,
        # even if below threshold. This avoids returning only generic always-apply rules.
        # But only if we're below minimum threshold
        # Apply minimum quality thresholds to avoid very low-scoring or semantically weak rules
        if len(results) < MIN_TOP_K:
            for sr in scored_rules:
                if len(results) >= MIN_TOP_K:
                    break
                if sr.rule.path not in seen_paths:
                    # Require minimum RRF score
                    if sr.score < FALLBACK_MIN_SCORE:
                        continue
                    # If embeddings available (indicated by non-zero embedding scores in results),
                    # require minimum embedding score
                    if self.use_embeddings and sr.embedding_score < FALLBACK_MIN_EMBEDDING:
                        continue
                    if "score" not in sr.selection_reasons:
                        sr.selection_reasons.append("score")
                    results.append(sr)
                    seen_paths.add(sr.rule.path)

        # Inject domain hub rules if requested (do not count toward threshold; mark injected in output).
        if inject_paths:
            for p in inject_paths:
                if p in seen_paths:
                    continue
                rule = self._store.get_rule_by_path(p)
                if rule is None:
                    continue
                results.append(
                    ScoredRule(
                        rule=rule,
                        score=0.0,
                        matched_tokens=set(),
                        injected=True,
                        selection_reasons=["injected"],
                    )
                )
                seen_paths.add(rule.path)

        # Expand with required dependencies
        if expand_required:
            results = self._expand_required_rules(results, seen_paths, scored_by_path)

        return results

    def _expand_required_rules(
        self,
        results: list[ScoredRule],
        seen_paths: set[str],
        scored_by_path: dict[str, ScoredRule],
    ) -> list[ScoredRule]:
        """Expand results to include required rules from retrieved rules.

        Only injects dependencies from rules with meaningful scores (above threshold or at least 0.15).
        This prevents injecting dependencies from low-scoring rules that were only included via fallback.
        """
        from aida_mcp.tools.rules.config import DEFAULT_SCORE_THRESHOLD

        # Minimum score for dependency injection (lower than threshold to allow some flexibility)
        min_score_for_injection = 0.15

        required_paths: dict[str, set[str]] = {}
        result_by_path = {sr.rule.path: sr for sr in results}
        for sr in results:
            # Only expand dependencies for rules that were actually selected by scoring.
            # This avoids dependency fan-out from always/injected-only entries.
            if "score" not in sr.selection_reasons:
                continue
            # Skip low-confidence scored entries (fallback noise guard).
            if sr.score < min(DEFAULT_SCORE_THRESHOLD, min_score_for_injection):
                continue
            for req_path in sr.rule.required:
                req_rule = self._store.get_rule_by_path(req_path)
                if req_rule is None:
                    continue
                # Additional guard: don't auto-include a required rule that is completely
                # unrelated to this query according to the scorer (common false positive case).
                req_scored = scored_by_path.get(req_rule.path)
                if req_scored is not None and req_scored.score < FALLBACK_MIN_SCORE:
                    continue

                if req_rule.path in seen_paths:
                    existing = result_by_path.get(req_rule.path)
                    if existing is not None:
                        if "required" not in existing.selection_reasons:
                            existing.selection_reasons.append("required")
                        if sr.rule.rule_id not in existing.required_by_rule_ids:
                            existing.required_by_rule_ids.append(sr.rule.rule_id)
                    continue

                required_paths.setdefault(req_rule.path, set()).add(sr.rule.rule_id)

        for req_path, required_by_ids in required_paths.items():
            rule = self._store.get_rule_by_path(req_path)
            if rule and rule.path not in seen_paths:
                results.append(
                    ScoredRule(
                        rule,
                        score=0.0,
                        matched_tokens=set(),
                        selection_reasons=["required"],
                        required_by_rule_ids=sorted(required_by_ids),
                    )
                )
                seen_paths.add(rule.path)

        return results
