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

from __future__ import annotations

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

import orjson
from pydantic import Field

from aida_mcp.output.policy import MAX_MCP_RESPONSE_BYTES
from aida_mcp.server.app_instance import Context
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


def register(mcp: FastMCP) -> None:
    """Register rules tools."""
    mcp.tool()(get_rules)
    mcp.tool()(warmup_embeddings)


async def get_rules(
    query: Annotated[
        str,
        Field(
            description=(
                "The user's question, task description, or conversation summary. "
                "For multi-turn conversations, summarize the overall goal/intent. "
                "Example: 'Adding gRPC endpoint to gen-ai service with tests' "
                "rather than just 'how do I add tests?'"
            )
        ),
    ],
    session_id: Annotated[
        str,
        Field(
            description=(
                "Mandatory per-chat session identifier used for server-side delta mode. "
                "Must be stable for the whole conversation, and unique across simultaneous chats "
                "(for example a UUID generated by the client per chat). "
                "With the same session_id, get_rules automatically remembers previous rule digests and "
                "returns only new/changed rule payloads when possible."
            )
        ),
    ],
    file_context: Annotated[
        list[str] | None,
        Field(
            description=(
                "Optional list of file paths currently being worked on. "
                "Helps retrieve rules relevant to specific services/components."
            )
        ),
    ] = None,
    workspace_root: Annotated[
        str | None,
        Field(
            description=(
                "Optional explicit workspace root path. If not provided, "
                "determined from MCP client roots or AIDA_WORKSPACE_ROOT env var."
            )
        ),
    ] = None,
    *,
    ctx: Context,
) -> RulesResponse:
    """
    Retrieve relevant development rules and guidelines based on conversation context.

    The returned rules are the authoritative, repo-specific guidance. This tool is intentionally designed to
    keep "how to do the task" in repo-owned rule files under `.aida/rules/**` (and required/related rules),
    rather than hardcoding repo rules into client-side instruction files.

    Expected usage (behavioral contract):
    - Call `get_rules` at the start of every task and again on scope change.
    - Query quality matters: pass a concrete goal/intent (one sentence is enough). Avoid vague queries.
    - Delta mode: pass a stable `session_id`; server keeps per-session known rule digests automatically.
    - Reporting: after calling `get_rules`, show a compact summary of:
      - selected rule IDs (from `rule_digests[*].rule_id`)
      - skipped rule IDs (from `skipped_rule_digests[*].rule_id`, when present)
      Full excerpts/scores are optional and should be shown only if the user asks or if debugging retrieval.
    - Artifact-first exception: if you are primarily analyzing external artifacts (PR/CI logs/tickets),
      fetch/read the artifact first, then call `get_rules` with the artifact details in the query.
    - Do not stall: after calling `get_rules` and printing the compact summary, continue with the task in the
      same response. If you need user input, ask a specific question and proceed with safe work in parallel.

    Skipping `get_rules` is acceptable only when clearly unnecessary:
    - follow-up questions on the same topic (rules already in context)
    - pure Q&A with no coding involved
    - discussing work already completed in the session
    """
    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.")

    # Determine workspace root: explicit param > MCP client roots > env var
    if workspace_root:
        resolved_root = Path(workspace_root)
    else:
        resolved_root = 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},
                    }
                ],
            },
        )

    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 = cast(RulesResponse, 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 explicit workspace root path. If not provided, "
                "determined from MCP client roots or AIDA_WORKSPACE_ROOT env var."
            )
        ),
    ] = None,
    *,
    ctx: Context,
) -> dict[str, object]:
    """
    Warm up semantic rule retrieval (download model + build embeddings index).

    This is optional. `get_rules` falls back to BM25-only retrieval if the embedding model isn't cached.
    Running this tool once per server start enables semantic scoring without surprising downloads during `get_rules`.
    """
    from aida_mcp.tools.rules.embeddings.initialization import initialize_embeddings_async
    from aida_mcp.tools.rules.materialize import materialize_rules_dir

    if workspace_root:
        resolved_root = Path(workspace_root)
    else:
        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,  # avoid progress token issues; this is an explicit warmup anyway
            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 _shrink_rules_response(response: RulesResponse) -> RulesResponse:
    """Ensure get_rules response stays under the MCP response size limit."""

    def _payload_size(payload: object, *, pretty: bool = False) -> int | None:
        try:
            if pretty:
                return len(orjson.dumps(payload, option=orjson.OPT_NON_STR_KEYS | orjson.OPT_INDENT_2))
            return len(orjson.dumps(payload, option=orjson.OPT_NON_STR_KEYS))
        except (TypeError, ValueError):
            return None

    size = _payload_size(response)
    size_pretty = _payload_size(response, pretty=True)
    if size is None or size_pretty is None:
        return response
    if size <= MAX_MCP_RESPONSE_BYTES and size_pretty <= MAX_MCP_RESPONSE_BYTES:
        return response

    if response.get("ok") is not True:
        return response

    rules_raw = response.get("rules", [])
    if not isinstance(rules_raw, list) or not rules_raw:
        minimal = _rules_payload_omitted()
        return cast(RulesResponse, minimal)

    rules_raw_list = cast(list[object], rules_raw)
    rules = [cast(dict[str, object], rule) for rule in rules_raw_list if isinstance(rule, dict)]
    if not rules:
        minimal = _rules_payload_omitted()
        return cast(RulesResponse, minimal)

    working = dict(response)
    remaining: list[dict[str, object]] = list(rules)
    while len(remaining) > 1:
        protected_ids: set[str] = _required_closure(remaining)
        removable: list[dict[str, object]] = [
            rule for rule in remaining if _rule_id(rule) not in protected_ids and not _rule_is_always_apply(rule)
        ]
        if not removable:
            break
        lowest: dict[str, object] = min(removable, key=_rule_score)
        remaining = [rule for rule in remaining if rule is not lowest]
        working["rules"] = remaining
        size = _payload_size(working)
        size_pretty = _payload_size(working, pretty=True)
        if size is None or size_pretty is None:
            break
        if size <= MAX_MCP_RESPONSE_BYTES and size_pretty <= MAX_MCP_RESPONSE_BYTES:
            remaining = _prune_orphan_dependencies(remaining)
            working["rules"] = remaining
            _sync_rule_metadata(working, remaining)
            return cast(RulesResponse, working)

    remaining = _prune_orphan_dependencies(remaining)
    trimmed = [dict(rule) for rule in remaining]
    for rule in trimmed:
        if "content" in rule:
            rule["content"] = {
                "kind": "omitted",
                "text": "Rule content omitted due to response size limits.",
            }
    working["rules"] = trimmed
    size = _payload_size(working)
    size_pretty = _payload_size(working, pretty=True)
    if (
        size is not None
        and size_pretty is not None
        and size <= MAX_MCP_RESPONSE_BYTES
        and size_pretty <= MAX_MCP_RESPONSE_BYTES
    ):
        _sync_rule_metadata(working, trimmed)
        return cast(RulesResponse, working)

    minimal = _rules_payload_omitted()
    return cast(RulesResponse, minimal)


def _sync_rule_metadata(payload: dict[str, object], rules: list[dict[str, object]]) -> None:
    rule_digests: list[dict[str, str]] = []
    selected_rule_ids: list[str] = []
    selected_rule_labels: list[str] = []
    for rule in rules:
        rule_id = rule.get("rule_id")
        sha256 = rule.get("sha256")
        if isinstance(rule_id, str):
            selected_rule_ids.append(rule_id)
            reasons = _rule_selection_reasons(rule)
            if reasons:
                selected_rule_labels.append(f"{rule_id}({'+'.join(reasons)})")
            else:
                selected_rule_labels.append(rule_id)
            if isinstance(sha256, str):
                rule_digests.append({"rule_id": rule_id, "sha256": sha256})

    payload["rule_digests"] = rule_digests
    metadata = payload.get("metadata")
    if isinstance(metadata, dict):
        md = cast(dict[str, object], metadata)
        md["rules_returned"] = len(rules)
        md["selected_rule_ids"] = selected_rule_ids
        md["selected_rule_labels"] = selected_rule_labels


def _rule_selection_reasons(rule: dict[str, object]) -> list[str]:
    value = rule.get("selection_reasons")
    if isinstance(value, list):
        reasons_raw = cast(list[object], value)
        reasons = [item for item in reasons_raw if isinstance(item, str)]
        ordered = [name for name in ("always", "score", "required", "injected") if name in reasons]
        if ordered:
            return ordered
    # Backward-compatible fallback if labels are not present in the payload.
    reasons: list[str] = []
    if rule.get("always_apply") is True:
        reasons.append("always")
    if _rule_score(rule) > 0.0:
        reasons.append("score")
    if rule.get("injected") is True:
        reasons.append("injected")
    return reasons


def _rules_payload_omitted() -> dict[str, object]:
    return {
        "ok": True,
        "metadata": None,
        "rule_digests": [],
        "rules": [],
        "message": "Rule payload omitted due to response size limits.",
    }


def _rule_id(rule: dict[str, object]) -> str | None:
    value = rule.get("rule_id")
    return value if isinstance(value, str) else None


def _rule_required_ids(rule: dict[str, object]) -> list[str]:
    value = rule.get("required")
    if not isinstance(value, list):
        return []
    required_raw = cast(list[object], value)
    return [item for item in required_raw if isinstance(item, str)]


def _rule_score(rule: dict[str, object]) -> float:
    raw = rule.get("score")
    if isinstance(raw, (int, float)):
        return float(raw)
    return float("-inf")


def _rule_is_always_apply(rule: dict[str, object]) -> bool:
    return bool(rule.get("always_apply") is True)


def _rule_is_primary(rule: dict[str, object]) -> bool:
    injected = rule.get("injected")
    always_apply = rule.get("always_apply")
    score = _rule_score(rule)
    return bool(injected is True or always_apply is True or score > 0.0)


def _required_closure(rules: list[dict[str, object]]) -> set[str]:
    by_id: dict[str, dict[str, object]] = {}
    for rule in rules:
        rid = _rule_id(rule)
        if rid:
            by_id[rid] = rule
    protected: set[str] = set()
    queue: list[str] = []
    for rule in rules:
        rid = _rule_id(rule)
        if rid and _rule_is_primary(rule):
            protected.add(rid)
            queue.append(rid)

    while queue:
        current = queue.pop()
        rule = by_id.get(current)
        if rule is None:
            continue
        for required_id in _rule_required_ids(rule):
            if required_id not in protected:
                protected.add(required_id)
                queue.append(required_id)
    return protected


def _prune_orphan_dependencies(rules: list[dict[str, object]]) -> list[dict[str, object]]:
    protected = _required_closure(rules)
    pruned: list[dict[str, object]] = []
    for rule in rules:
        if _rule_is_always_apply(rule):
            pruned.append(rule)
            continue
        rid = _rule_id(rule)
        if rid and rid in protected:
            pruned.append(rule)
    return pruned
