# Copyright (C) 2026 GoodData Corporation

"""Payload-size reduction helpers for rules tool responses."""

from __future__ import annotations

from collections.abc import Callable
from typing import cast

import orjson

from aida_mcp.output.policy import MAX_MCP_RESPONSE_BYTES
from aida_mcp.tools.rules.formatting import RulesResponse


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

    if "context" in response:
        return cast(RulesResponse, _shrink_compact_response(cast(dict[str, object], response), _payload_size))
    if "rules" in response:
        return cast(RulesResponse, _shrink_debug_response(cast(dict[str, object], response), _payload_size))
    return response


def _shrink_compact_response(
    response: dict[str, object],
    payload_size: Callable[..., int | None],
) -> dict[str, object]:
    context = response.get("context")
    reporting = response.get("reporting")
    if not isinstance(context, dict) or not isinstance(reporting, dict):
        return response
    context_obj = cast(dict[str, object], context)
    reporting_obj = cast(dict[str, object], reporting)
    rules_raw = context_obj.get("rules")
    selected_raw = reporting_obj.get("selected")
    if not isinstance(rules_raw, list) or not isinstance(selected_raw, list):
        return response

    rules = [cast(dict[str, object], rule) for rule in rules_raw if isinstance(rule, dict)]
    selected = [cast(dict[str, object], item) for item in selected_raw if isinstance(item, dict)]
    if not rules:
        return response

    working = dict(response)
    remaining_rules = list(rules)
    remaining_selected = list(selected)
    while len(remaining_rules) > 1:
        removable_indexes = [idx for idx, item in enumerate(remaining_selected) if not _compact_rule_is_protected(item)]
        if not removable_indexes:
            break
        remove_idx = removable_indexes[-1]
        remaining_rules.pop(remove_idx)
        remaining_selected.pop(remove_idx)
        _set_compact_payload(working, remaining_rules, remaining_selected)
        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
        ):
            return working

    trimmed_rules = [dict(rule) for rule in remaining_rules]
    for rule in trimmed_rules:
        rule["content"] = {
            "kind": "omitted",
            "text": "Rule content omitted due to response size limits.",
        }
    _set_compact_payload(working, trimmed_rules, remaining_selected)
    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
    ):
        return working

    return {
        "ok": True,
        "context": {"low_confidence": False, "rules": []},
        "reporting": {
            "required": True,
            "guidance": "Rule payload omitted due to response size limits.",
            "selected": [],
        },
        "message": "Rule payload omitted due to response size limits.",
    }


def _set_compact_payload(
    payload: dict[str, object],
    rules: list[dict[str, object]],
    selected: list[dict[str, object]],
) -> None:
    context = cast(dict[str, object], payload["context"])
    reporting = cast(dict[str, object], payload["reporting"])
    context["rules"] = rules
    reporting["selected"] = selected


def _compact_rule_is_protected(item: dict[str, object]) -> bool:
    reasons = item.get("reasons")
    if not isinstance(reasons, list):
        return False
    reason_values = {reason for reason in reasons if isinstance(reason, str)}
    return bool({"always", "required", "injected"} & reason_values)


def _shrink_debug_response(
    response: dict[str, object],
    payload_size: Callable[..., int | None],
) -> dict[str, object]:
    rules_raw = response.get("rules", [])
    if not isinstance(rules_raw, list) or not rules_raw:
        return {
            "ok": True,
            "metadata": None,
            "rule_digests": [],
            "rules": [],
            "message": "Rule payload omitted due to response size limits.",
        }
    rules = [cast(dict[str, object], rule) for rule in rules_raw if isinstance(rule, dict)]
    if not rules:
        return {
            "ok": True,
            "metadata": None,
            "rule_digests": [],
            "rules": [],
            "message": "Rule payload omitted due to response size limits.",
        }

    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 not None
            and size_pretty is not None
            and 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 working

    remaining = _prune_orphan_dependencies(remaining)
    trimmed = [dict(rule) for rule in remaining]
    for rule in trimmed:
        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 working

    return {
        "ok": True,
        "metadata": None,
        "rule_digests": [],
        "rules": [],
        "message": "Rule payload omitted due to response size limits.",
    }


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 = [item for item in value if isinstance(item, str)]
        ordered = [name for name in ("always", "score", "required", "injected") if name in reasons]
        if ordered:
            return ordered
    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 _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 []
    return [item for item in value 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
