# Copyright (C) 2026 GoodData Corporation

"""Rules retrieval tool for MCP-based rule delivery.

Main entry point for rule retrieval. Implementation details are in aida_mcp.tools.rules.*
"""

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING, Any, cast

from aida_mcp.tools.rules.config import DEFAULT_SCORE_THRESHOLD
from aida_mcp.tools.rules.embeddings.initialization import (
    get_vector_store,
    initialize_embeddings,
    initialize_embeddings_async,
)
from aida_mcp.tools.rules.formatting import (
    RetrievalMetadata,
    RuleDigest,
    RulesResponse,
    RulesResponseDebugOk,
    format_rules_response,
    ordered_selection_reasons,
)
from aida_mcp.tools.rules.injection import compute_inject_paths
from aida_mcp.tools.rules.materialize import materialize_rules_dir
from aida_mcp.tools.rules.retriever import RulesRetriever
from aida_mcp.tools.rules.scoring import ScoredRule
from aida_mcp.tools.rules.utils import enhance_query_with_file_context

if TYPE_CHECKING:
    from mcp.server.fastmcp import Context as MCPContext

    Context = MCPContext[Any, Any, Any]

__all__ = [
    "get_relevant_rules",
    "get_vector_store",
    "initialize_embeddings",
    "RulesRetriever",
]


def _rule_label(sr: ScoredRule) -> str:
    reasons = ordered_selection_reasons(sr)
    return f"{sr.rule.rule_id}({'+'.join(reasons)})" if reasons else sr.rule.rule_id


async def get_relevant_rules(
    query: str,
    workspace_root: Path,
    file_context: list[str] | None = None,
    known_rule_digests: list[RuleDigest] | None = None,
    *,
    debug: bool = False,
    ctx: Context | None = None,
) -> RulesResponse:
    """
    Main entry point for rule retrieval.

    Args:
        query: The user's question or task description
        workspace_root: Path to the workspace root
        file_context: Optional list of file paths for additional context
        known_rule_digests: Optional list of known rule digests (for delta mode)
        debug: Whether to return retrieval diagnostics and full debug payload
        ctx: Optional MCP context for progress reporting

    Returns:
        JSON-compatible dict with compact context/reporting or full debug payload
    """
    # Materialize rule sources into a single directory for retrieval:
    # - embedded rules shipped with aida-mcp
    # - repo-owned `.aida/rules/`
    rules_dir = materialize_rules_dir(workspace_root=workspace_root)

    # Best-effort semantic index initialization.
    #
    # By default we do NOT download the embedding model as a side effect of get_rules.
    # If the model is already cached locally, we'll build the in-memory vector store and
    # enable semantic scoring. Otherwise, retrieval falls back to BM25-only.
    semantic_available = await initialize_embeddings_async(rules_dir, ctx, allow_download=False, suppress_errors=True)

    # Embeddings must be warmed up explicitly. If unavailable, fail fast with guidance.
    #
    # Rationale: we want deterministic behavior and to avoid surprise network calls,
    # but also require hybrid retrieval quality once the repo is AIDA-initialized.
    if not semantic_available:
        raise RuntimeError(
            "Embeddings are not available for `get_rules` (the embedding model is not cached).\n\n"
            "Fix:\n"
            "- call MCP tool: `warmup_embeddings` (optionally pass workspace_root)\n"
            "- then retry `get_rules`\n\n"
            "Why: `get_rules` does not download models as a side effect; warmup performs the one-time download."
        )

    # Enhance query with file context if provided
    enhanced_query = enhance_query_with_file_context(query, file_context)

    inject_paths = compute_inject_paths(query=query, file_context=file_context)

    # Create retriever and get results
    # Note: Required dependencies are automatically expanded via rule frontmatter
    retriever = RulesRetriever(rules_dir, vector_store=get_vector_store(rules_dir))
    scored_rules = retriever.retrieve(enhanced_query, inject_paths=inject_paths or None)

    # Optional delta filtering: if caller already "knows" some rules (rule_id+sha256), skip returning their full payload.
    known: set[tuple[str, str]] = set()
    skipped: list[ScoredRule] = []
    if known_rule_digests:
        for d in known_rule_digests:
            rid = d.get("rule_id")
            sha = d.get("sha256")
            if rid and sha:
                known.add((rid, sha))

    if known:
        kept: list[ScoredRule] = []
        for sr in scored_rules:
            if (sr.rule.rule_id, sr.rule.rule_hash) in known:
                skipped.append(sr)
            else:
                kept.append(sr)
        scored_rules = kept

    # Build metadata
    non_always = [sr for sr in scored_rules if not sr.rule.always_apply]
    low_confidence = bool(non_always) and all(
        sr.score < DEFAULT_SCORE_THRESHOLD for sr in non_always if not sr.injected
    )
    bm25_only_hint: str | None = None
    metadata = RetrievalMetadata(
        query=query,
        score_threshold=DEFAULT_SCORE_THRESHOLD,
        rules_searched=len(retriever.store.get_all_rules()),
        rules_returned=len(scored_rules),
        low_confidence=low_confidence,
        selected_rule_ids=[sr.rule.rule_id for sr in scored_rules],
        selected_rule_labels=[_rule_label(sr) for sr in scored_rules],
        skipped_rule_ids=[sr.rule.rule_id for sr in skipped] if skipped else [],
        reminder=(
            "Required: report selected/skipped rules after get_rules. "
            "Prefer metadata.selected_rule_labels for provenance "
            "(always/score/required/injected), e.g. rule-id(score+required)."
        ),
        semantic_available=semantic_available,
        bm25_only_hint=bm25_only_hint,
    )

    resp = format_rules_response(scored_rules, metadata, debug=debug)
    if debug and resp.get("ok") is True and known:
        debug_resp = cast(RulesResponseDebugOk, resp)
        debug_resp["skipped_rule_digests"] = cast(
            list[RuleDigest],
            [{"rule_id": sr.rule.rule_id, "sha256": sr.rule.rule_hash} for sr in skipped],
        )
    return resp
