# 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)


def _mcp_session_id_from_context(ctx: Context) -> str | None:
    """Return the MCP transport session ID when available."""
    request_context = getattr(ctx, "request_context", None)
    request = getattr(request_context, "request", None) if request_context is not None else None
    headers = getattr(request, "headers", None)

    if headers is not None:
        try:
            return headers.get("mcp-session-id") or headers.get("Mcp-Session-Id")
        except Exception:  # noqa: BLE001
            return None
    return None


def _resolve_effective_session_id(explicit_session_id: str, ctx: Context) -> tuple[str, str]:
    """Resolve the effective rules session identifier and its source."""
    mcp_session_id = _mcp_session_id_from_context(ctx)
    if isinstance(mcp_session_id, str) and mcp_session_id.strip():
        return mcp_session_id.strip(), "mcp_session_id"
    return explicit_session_id, "explicit_session_id"


def _transport_debug_info(ctx: Context, *, effective_session_source: str) -> dict[str, object]:
    """Return minimal transport/session hints for debug-only inspection."""
    request_context = getattr(ctx, "request_context", None)
    request = getattr(request_context, "request", None) if request_context is not None else None
    headers = getattr(request, "headers", None)

    return {
        "client_id": getattr(ctx, "client_id", None),
        "request_id": getattr(ctx, "request_id", None),
        "has_request_context": request_context is not None,
        "request_type": type(request).__name__ if request is not None else None,
        "has_headers": headers is not None,
        "mcp_session_id": _mcp_session_id_from_context(ctx),
        "effective_session_source": effective_session_source,
    }


def _delta_empty_response(*, merged_cursor: str) -> RulesResponse:
    return {
        "ok": True,
        "delta": {"cursor": merged_cursor, "mode": "delta_empty"},
        "message": "No new or changed rules.",
    }


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)."
            )
        ),
    ],
    delta_cursor: Annotated[
        str | None,
        Field(
            description=(
                "Opaque cursor from the previous get_rules response. Clients should persist the latest cursor "
                "and pass it on every subsequent get_rules call in the same conversation while the earlier "
                "rule payload is still retained in context. Omitting this argument disables delta mode and "
                "forces a full response."
            )
        ),
    ] = None,
    file_context: Annotated[
        list[str] | None,
        Field(
            description=(
                "Optional list of file paths in scope. Improves rule relevance. "
                "Pass a JSON array/list of repo-relative paths such as "
                '["packages/aida-mcp/src/aida_mcp/server/tools_rules.py", "packages/aida-mcp/tests/test_tools_rules.py"]. '
                "Do not pass a single string."
            )
        ),
    ] = None,
    debug: Annotated[
        bool,
        Field(
            description=(
                "When true, return retrieval diagnostics and the full legacy debug payload. "
                "Defaults to false for a compact context/reporting response."
            )
        ),
    ] = False,
    *,
    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` as a list of active file paths, for example
      `file_context=[\"packages/aida-mcp/src/aida_mcp/server/tools_rules.py\"]`.
    - Use `debug=true` only when troubleshooting retrieval behavior or inspecting internals.

    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.
    - Normal mode returns two use-case sections: `context` for rule content and `reporting`
      for concise provenance to surface back to the user.
    """
    from aida_mcp.tools.rules.session_memory import (
        get_session_cursor,
        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.")
    effective_session_id, effective_session_source = _resolve_effective_session_id(normalized_session_id, ctx)

    resolved_root = await get_workspace_root(ctx)
    session_key = f"{resolved_root.resolve()}::{effective_session_id}"
    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()
        bootstrap = cast(
            RulesResponse,
            _bootstrap_response(
                query=query,
                workspace_root=resolved_root,
                config_root=readiness.config_root,
                missing_validate_files=list(readiness.missing_validate_files),
                text=text,
                sha=sha,
                debug=debug,
            ),
        )
        _, merged_cursor = merge_session_rule_digests(
            session_key,
            [{"rule_id": "general/aida-bootstrap.mdc", "sha256": sha}],
        )
        if isinstance(bootstrap, dict):
            bootstrap_payload = cast(dict[str, object], bootstrap)
            bootstrap_payload["delta"] = {"cursor": merged_cursor}
        return bootstrap

    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."
                ),
            },
        )

    current_cursor = get_session_cursor(session_key)
    known = get_session_known_rule_digests(session_key) if delta_cursor and delta_cursor == current_cursor else []
    used_delta = bool(known)
    response = await get_relevant_rules(query, resolved_root, file_context, known, debug=debug, ctx=ctx)
    shrunk = shrink_rules_response(response)
    if shrunk.get("ok") is True:
        digests = response.get("rule_digests", []) if isinstance(response, dict) else []
        known_count, merged_cursor = merge_session_rule_digests(session_key, digests)
        if isinstance(shrunk, dict):
            shrunk_payload = cast(dict[str, object], shrunk)
            delta_mode = "delta_empty" if used_delta and not digests else "full"
            if delta_mode == "delta_empty":
                return _delta_empty_response(merged_cursor=merged_cursor)
            shrunk_payload["delta"] = {"cursor": merged_cursor, "mode": delta_mode}
            if not debug:
                shrunk_payload.pop("rule_digests", None)
            reporting = shrunk_payload.get("reporting")
            if isinstance(reporting, dict):
                reporting_payload = cast(dict[str, object], reporting)
                reporting_payload["guidance"] = (
                    "Required: report only selected rule ids with reasons, e.g. "
                    "`core/tools/workflow-router.mdc(always)`. "
                    "Report `low_confidence: true` only when true."
                )
        metadata = shrunk.get("metadata") if isinstance(shrunk, dict) else None
        if isinstance(metadata, dict):
            md = cast(dict[str, object], metadata)
            md["session_id"] = normalized_session_id
            md["effective_session_id"] = effective_session_id
            md["known_rule_digests_count"] = known_count
            md["delta_cursor"] = merged_cursor
            md["response_mode"] = "delta_empty" if used_delta and not digests else "full"
            if debug:
                md["transport_debug"] = _transport_debug_info(ctx, effective_session_source=effective_session_source)
    return shrunk


def _bootstrap_response(
    *,
    query: str,
    workspace_root: Path,
    config_root: Path,
    missing_validate_files: list[str],
    text: str,
    sha: str,
    debug: bool,
) -> dict[str, object]:
    if debug:
        return {
            "ok": True,
            "metadata": {
                "query": query,
                "repo_initialized": False,
                "workspace_root": str(workspace_root),
                "config_root": str(config_root),
                "missing_validate_files": 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": [],
                    "origin": "system",
                    "editable": False,
                    "content": {"kind": "full", "text": text},
                }
            ],
        }
    return {
        "ok": True,
        "context": {
            "low_confidence": False,
            "rules": [
                {
                    "id": "general/aida-bootstrap.mdc",
                    "origin": "system",
                    "editable": False,
                    "content": {"kind": "full", "text": text},
                }
            ],
        },
        "reporting": {
            "required": True,
            "guidance": "Report selected rules concisely with provenance (id, origin, selection reasons).",
            "selected": [
                {
                    "id": "general/aida-bootstrap.mdc",
                    "origin": "system",
                    "editable": False,
                    "reasons": ["injected"],
                }
            ],
        },
    }


async def warmup_embeddings(*, 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 = 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"]
