from __future__ import annotations

import fnmatch
from typing import Any, Iterable

from pydantic_ai.exceptions import ModelRetry
from pydantic_ai.toolsets.wrapper import WrapperToolset

from ..audit import AuditEvent, tags_to_strings
from ..config import get_config
from ..context import get_current_context
from ..decorators import apply_directives
from ..exceptions import ApprovalRequired
from ..governance import evaluate_policies, is_visible, record_policy_event, record_visibility_event
from ..types import PolicyRequest, ToolRules, as_frozenset


def _visibility_block_message(rules: ToolRules, tool_name: str) -> str:
    reason = "blocked_by" if rules.blocked_by else "boundary"
    boundary = str(rules.boundary) if rules.boundary is not None else "unknown"
    if reason == "blocked_by":
        return f"Tool '{tool_name}' is blocked by active tags."
    return f"Tool '{tool_name}' is not available due to boundary '{boundary}'."

class RegulatedMCPAdapter(WrapperToolset[Any]):
    def __init__(
        self,
        *,
        wrapped: Any | None = None,
        server: Any | None = None,
        tools: dict[str, bool | list[Any]] | None = None,
        activates: Iterable[Any] | None = None,
        blocked_by: Iterable[Any] | None = None,
        boundary: Any | None = None,
    ) -> None:
        if wrapped is None:
            wrapped = server
        if wrapped is None:
            raise ValueError("RegulatedMCPAdapter requires a wrapped MCP server")
        super().__init__(wrapped=wrapped)
        self._tools = tools or {}
        self._server_rules = ToolRules(
            activates=as_frozenset(activates),
            blocked_by=as_frozenset(blocked_by),
            boundary=boundary,
        )
        self._last_visible_regulated: set[str] = set()
        self._last_visible_unregulated: set[str] = set()
        self._last_hidden: set[str] = set()
        self._last_filtered: set[str] = set()

    def _resolve_config(self, tool_name: str) -> bool | list[Any] | None:
        if not self._tools:
            return True
        if tool_name in self._tools:
            return self._tools[tool_name]
        matches = [
            (pattern, config)
            for pattern, config in self._tools.items()
            if fnmatch.fnmatch(tool_name, pattern)
        ]
        if not matches:
            return None

        def specificity(item: tuple[str, Any]) -> tuple[int, int]:
            pattern = item[0]
            return (pattern.count("*"), -len(pattern))

        matches.sort(key=specificity)
        return matches[0][1]

    def _server_label(self) -> str | None:
        return getattr(self.wrapped, "id", None) or getattr(self.wrapped, "label", None) or getattr(self, "label", None)

    def _rules_for(self, tool_name: str) -> ToolRules | None:
        config = self._resolve_config(tool_name)
        if config is None or config is False:
            return None
        if config is True:
            return self._server_rules
        if isinstance(config, list):
            return apply_directives(self._server_rules, config)
        raise TypeError(f"Invalid MCP tool config for {tool_name!r}: {config!r}")

    def rules_for_tool(self, name: str) -> ToolRules | None:
        return self._rules_for(name)

    def coverage_summary(self) -> dict[str, Any]:
        tools_config: dict[str, Any] = {}
        for pattern, config in self._tools.items():
            if config is False or config is True:
                tools_config[pattern] = config
                continue
            if isinstance(config, list):
                rules = apply_directives(self._server_rules, config)
                tools_config[pattern] = _rules_payload(rules)
                continue
            raise TypeError(f"Invalid MCP tool config for {pattern!r}: {config!r}")

        return {
            "tools": tools_config,
            "server_rules": _rules_payload(self._server_rules),
            "active_tools": {
                "regulated": sorted(self._last_visible_regulated),
                "non_regulated": sorted(self._last_visible_unregulated),
            },
            "hidden_tools": sorted(self._last_hidden),
            "filtered_tools": sorted(self._last_filtered),
        }

    async def get_tools(self, ctx: Any) -> dict[str, Any]:
        tools = await self.wrapped.get_tools(ctx)
        config = get_config()
        current = get_current_context()
        active_tags = current.active_tags if current else set()
        filtered: dict[str, Any] = {}
        visible_regulated: set[str] = set()
        visible_unregulated: set[str] = set()
        hidden: set[str] = set()
        filtered_out: set[str] = set()
        for name, tool in tools.items():
            rules = self._rules_for(name)
            if rules is None:
                filtered_out.add(name)
                continue
            visible = is_visible(rules, active_tags, config.boundaries)
            record_visibility_event(
                current,
                name,
                rules,
                active_tags,
                visible,
                tool_kind="mcp",
                mcp_server=self._server_label(),
            )
            if visible:
                filtered[name] = tool
                if rules.is_empty():
                    visible_unregulated.add(name)
                else:
                    visible_regulated.add(name)
            else:
                hidden.add(name)
        self._last_visible_regulated = visible_regulated
        self._last_visible_unregulated = visible_unregulated
        self._last_hidden = hidden
        self._last_filtered = filtered_out
        if current is not None:
            current.audit(
                AuditEvent(
                    event_type="tools_presented",
                    tool_kind="mcp",
                    tool_name=None,
                    mcp_server=self._server_label(),
                    agent_id=current.agent_id,
                    run_id=current.run_id,
                    active_tags=tags_to_strings(active_tags),
                    payload={"tools": sorted(filtered.keys())},
                )
            )
        return filtered

    async def call_tool(self, name: str, tool_args: dict[str, Any], ctx: Any, tool: Any) -> Any:
        rules = self._rules_for(name)
        if rules is None:
            raise ModelRetry(f"Unknown tool name: {name!r}.")

        current = get_current_context()
        pre_tags = set(current.active_tags) if current else set()
        tool_call_id = getattr(ctx, "tool_call_id", None)
        if current is not None:
            current.audit(
                AuditEvent(
                    event_type="tool_call_requested",
                    tool_kind="mcp",
                    tool_name=name,
                    tool_call_id=tool_call_id,
                    mcp_server=self._server_label(),
                    agent_id=current.agent_id,
                    run_id=current.run_id,
                    active_tags=tags_to_strings(pre_tags),
                )
            )
        visible = is_visible(rules, pre_tags, get_config().boundaries)
        if not visible:
            record_visibility_event(
                current,
                name,
                rules,
                pre_tags,
                visible,
                tool_kind="mcp",
                mcp_server=self._server_label(),
            )
            raise ModelRetry(_visibility_block_message(rules, name))

        decision = await evaluate_policies(
            rules,
            PolicyRequest(
                tool_name=name,
                args=tool_args,
                active_tags=pre_tags,
                boundary=rules.boundary,
                agent_id=current.agent_id if current else None,
                run_id=current.run_id if current else "no_context",
                run_context=ctx,
            ),
        )
        if decision.outcome != "allow":
            record_policy_event(
                current,
                name,
                rules,
                decision,
                tool_kind="mcp",
                enforced=get_config().mode != "monitor",
                mcp_server=self._server_label(),
                tool_call_id=tool_call_id,
            )
            if get_config().mode != "monitor":
                if decision.outcome == "approval":
                    raise ApprovalRequired(metadata={"reason": decision.reason})
                raise ModelRetry(decision.reason)

        would_block = get_config().mode == "monitor" and decision.outcome == "deny"
        would_require_approval = get_config().mode == "monitor" and decision.outcome == "approval"
        result = await self.wrapped.call_tool(name, tool_args, ctx, tool)

        if current is not None and rules.activates:
            current.activate(rules.activates)

        if current is not None:
            current.audit(
                AuditEvent(
                    event_type="tool_allowed",
                    tool_kind="mcp",
                    tool_name=name,
                    tool_call_id=tool_call_id,
                    mcp_server=self._server_label(),
                    agent_id=current.agent_id,
                    run_id=current.run_id,
                    active_tags=tags_to_strings(pre_tags),
                    active_tags_after=tags_to_strings(current.active_tags),
                    would_block=would_block if get_config().mode == "monitor" else None,
                    would_require_approval=would_require_approval if get_config().mode == "monitor" else None,
                )
            )

        return result


def _rules_payload(rules: ToolRules) -> dict[str, Any]:
    def _policy_names(values: Iterable[Any]) -> list[str]:
        names: list[str] = []
        for policy in values:
            name = getattr(policy, "__name__", None)
            names.append(name or repr(policy))
        return sorted(set(names))

    return {
        "boundary": str(rules.boundary) if rules.boundary is not None else None,
        "policies": _policy_names([*rules.policies.require, *rules.policies.any_of]),
        "tag": {
            "activates": [getattr(tag, "value", str(tag)) for tag in rules.activates],
            "blocked_by": [getattr(tag, "value", str(tag)) for tag in rules.blocked_by],
        },
    }
