# Copyright (C) 2026 GoodData Corporation

"""Rules-retrieval MCP tools."""

from __future__ import annotations

import hashlib
from pathlib import Path
from typing import TYPE_CHECKING, Annotated, cast

from pydantic import Field

from aida_mcp.server.app_instance import Context
from aida_mcp.server.tools_rules_payload import _sync_rule_metadata, shrink_rules_response
from aida_mcp.tools.rules.formatting import RulesResponse
from aida_mcp.utils.repo_readiness import check_repo_readiness
from aida_mcp.utils.workspace import get_workspace_root

if TYPE_CHECKING:
    from mcp.server.fastmcp import FastMCP

ONBOARDING_QUERY_TERMS: tuple[str, ...] = (
    "onboarding",
    "finalize onboarding",
    "aida-mcp init",
    "aida-mcp migrate",
    "aida-mcp doctor",
    "migrate rules",
)


def register(mcp: FastMCP) -> None:
    mcp.tool()(get_rules)
    mcp.tool()(warmup_embeddings)


async def get_rules(
    query: Annotated[
        str,
        Field(description=("Task goal in one sentence. Use concrete intent (what you are changing and where).")),
    ],
    session_id: Annotated[
        str,
        Field(
            description=(
                "Required stable ID for one chat. Enables delta mode (return only new/changed rules across calls)."
            )
        ),
    ],
    file_context: Annotated[
        list[str] | None,
        Field(description=("Optional file paths in scope. Improves rule relevance.")),
    ] = None,
    workspace_root: Annotated[
        str | None,
        Field(description=("Optional workspace root. If omitted, resolved from MCP roots/env.")),
    ] = None,
    *,
    ctx: Context,
) -> RulesResponse:
    """
    Retrieve task-relevant rules from embedded + repository rule sets.

    When to use:
    - At task start.
    - Again when scope changes significantly.

    How to use:
    - Provide a concrete `query`.
    - Keep `session_id` stable during the chat for delta mode.
    - Optionally pass `file_context` with active files.

    Notes:
    - If this tool errors asking for `warmup_embeddings`, run that tool once and retry.
    - For uninitialized repos, returns bootstrap guidance instead of normal retrieval.
    - For partially initialized repos, returns onboarding_incomplete unless this call is onboarding-focused.
    """
    from aida_mcp.tools.rules.session_memory import get_session_known_rule_digests, merge_session_rule_digests
    from aida_mcp.tools.rules_retrieval import get_relevant_rules

    await ctx.info(f"📚 Retrieving rules for: {query[:50]}...")
    normalized_session_id = session_id.strip()
    if not normalized_session_id:
        raise ValueError("session_id must be a non-empty string.")

    resolved_root = Path(workspace_root) if workspace_root else await get_workspace_root(ctx)
    readiness = check_repo_readiness(workspace_root=resolved_root)
    if not readiness.rules_ready:
        text = (
            "This repository is not initialized for AIDA.\n\n"
            "To enable repo-owned validation + rules:\n"
            "- create `.aida/` in the repo root (or set `AIDA_CONFIG_ROOT`)\n"
            "- add `.aida/change_domains.yaml`, `.aida/validation_policy.yaml`, `.aida/validation_registry.yaml`\n"
            "- optionally add repo rules under `.aida/rules/`\n\n"
            "Recommended: run `aida-mcp init`."
        )
        sha = hashlib.sha256(text.encode("utf-8")).hexdigest()
        return cast(
            RulesResponse,
            {
                "ok": True,
                "metadata": {
                    "query": query,
                    "repo_initialized": False,
                    "workspace_root": str(resolved_root),
                    "config_root": str(readiness.config_root),
                    "missing_validate_files": list(readiness.missing_validate_files),
                },
                "rule_digests": [{"rule_id": "general/aida-bootstrap.mdc", "sha256": sha}],
                "rules": [
                    {
                        "path": "<bootstrap>",
                        "rule_id": "general/aida-bootstrap.mdc",
                        "sha256": sha,
                        "description": "Bootstrap: initialize AIDA in this repository",
                        "always_apply": True,
                        "required": [],
                        "related": [],
                        "injected": True,
                        "score": 1.0,
                        "bm25_score": 0.0,
                        "embedding_score": 0.0,
                        "matched_tokens": [],
                        "selection_reasons": ["injected"],
                        "required_by_rule_ids": [],
                        "content": {"kind": "full", "text": text},
                    }
                ],
            },
        )

    onboarding_reasons: list[str] = []
    if readiness.missing_validate_files:
        onboarding_reasons.append("missing_validate_files")
    if not readiness.rules_dir_exists:
        onboarding_reasons.append("missing_rules_dir")
    elif not _rules_dir_has_mdc_rules(readiness.config_root):
        onboarding_reasons.append("empty_rules_dir")

    if onboarding_reasons and not _is_onboarding_query(query):
        missing_text = ", ".join(onboarding_reasons)
        return cast(
            RulesResponse,
            {
                "ok": False,
                "error": "onboarding_incomplete",
                "message": (
                    "AIDA onboarding is incomplete and rules retrieval is blocked for normal tasks.\n\n"
                    f"Detected reasons: {missing_text}\n\n"
                    "Call get_rules again with an onboarding-focused query to finalize setup, for example:\n"
                    '- query: "finalize onboarding after aida-mcp migrate"\n'
                    "- session_id: keep the same current session identifier."
                ),
            },
        )

    session_key = f"{resolved_root.resolve()}::{normalized_session_id}"
    known = get_session_known_rule_digests(session_key)
    response = await get_relevant_rules(query, resolved_root, file_context, known, ctx)
    shrunk = shrink_rules_response(response)
    if shrunk.get("ok") is True:
        ok = shrunk
        digests = ok.get("rule_digests", [])
        known_count = merge_session_rule_digests(session_key, digests)
        metadata = ok.get("metadata")
        if isinstance(metadata, dict):
            md = cast(dict[str, object], metadata)
            md["session_id"] = normalized_session_id
            md["known_rule_digests_count"] = known_count
    return shrunk


async def warmup_embeddings(
    workspace_root: Annotated[
        str | None,
        Field(description=("Optional workspace root. If omitted, resolved from MCP roots/env.")),
    ] = None,
    *,
    ctx: Context,
) -> dict[str, object]:
    """
    Download the embedding model and build the semantic index for rule retrieval.

    DO NOT call this tool proactively. Only call it when `get_rules` returns an error
    asking you to run `warmup_embeddings`. The `get_rules` tool auto-initializes from the
    local model cache; this tool is only needed for the one-time initial model download.
    """
    from aida_mcp.tools.rules.embeddings.initialization import initialize_embeddings_async
    from aida_mcp.tools.rules.materialize import materialize_rules_dir

    resolved_root = Path(workspace_root) if workspace_root else await get_workspace_root(ctx)
    rules_dir = materialize_rules_dir(workspace_root=resolved_root)
    await ctx.info("🤖 Warming up semantic rules index (may download model on first run)...")
    try:
        ok = await initialize_embeddings_async(
            rules_dir,
            ctx=None,
            allow_download=True,
            suppress_errors=False,
        )
    except Exception as e:
        return {
            "tool": "warmup_embeddings",
            "success": False,
            "text": f"Warmup failed: {type(e).__name__}: {e}",
            "workspace_root": str(resolved_root),
            "rules_dir": str(rules_dir),
        }

    if ok:
        return {
            "tool": "warmup_embeddings",
            "success": True,
            "text": "Warmup completed: semantic scoring enabled for get_rules.",
            "workspace_root": str(resolved_root),
            "rules_dir": str(rules_dir),
        }
    return {
        "tool": "warmup_embeddings",
        "success": True,
        "text": "Warmup skipped: embedding model not available (BM25-only).",
        "workspace_root": str(resolved_root),
        "rules_dir": str(rules_dir),
    }


def _is_onboarding_query(query: str) -> bool:
    q = query.lower()
    return any(term in q for term in ONBOARDING_QUERY_TERMS)


def _rules_dir_has_mdc_rules(config_root: Path) -> bool:
    rules_dir = config_root / "rules"
    if not rules_dir.exists() or not rules_dir.is_dir():
        return False
    return any(rules_dir.rglob("*.mdc"))


__all__ = ["_sync_rule_metadata", "get_rules", "register", "warmup_embeddings"]
