from __future__ import annotations

from collections.abc import Awaitable, Callable, Sequence
from typing import Any

from pydantic_ai.exceptions import ModelRetry, UserError
from pydantic_ai.toolsets.function import FunctionToolset
from pydantic_ai.toolsets.wrapper import WrapperToolset
from pydantic_ai.tools import Tool

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


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 RegulatedToolset(WrapperToolset[Any]):
    def __init__(
        self,
        tools: Sequence[Callable[..., Awaitable[Any]] | Tool[Any]] | None = None,
        *,
        toolset: FunctionToolset[Any] | None = None,
        id: str | None = None,
        max_retries: int = 1,
        timeout: float | None = None,
        docstring_format: str = "auto",
        require_parameter_descriptions: bool = False,
        schema_generator: type | None = None,
        strict: bool | None = None,
        sequential: bool = False,
        requires_approval: bool = False,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        if toolset is not None and tools is not None:
            raise UserError("Provide either tools or toolset, not both")
        if toolset is None:
            toolset_kwargs: dict[str, Any] = {
                "tools": list(tools or ()),
                "id": id,
                "max_retries": max_retries,
                "timeout": timeout,
                "docstring_format": docstring_format,
                "require_parameter_descriptions": require_parameter_descriptions,
                "strict": strict,
                "sequential": sequential,
                "requires_approval": requires_approval,
                "metadata": metadata,
            }
            if schema_generator is not None:
                toolset_kwargs["schema_generator"] = schema_generator
            toolset = FunctionToolset(**toolset_kwargs)
        super().__init__(wrapped=toolset)
        self._rules_by_name: dict[str, ToolRules] = {}
        for item in toolset.tools.values():
            tool_name = item.name
            func = item.function
            self._rules_by_name[tool_name] = rules_for_tool(func)

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

    def tool_rules_by_name(self) -> dict[str, ToolRules]:
        return dict(self._rules_by_name)

    def visit_and_replace(self, visitor: Callable[[Any], Any]) -> Any:
        wrapped = self.wrapped.visit_and_replace(visitor)
        if wrapped is self.wrapped:
            return self
        if not isinstance(wrapped, FunctionToolset):
            return self
        return RegulatedToolset(toolset=wrapped)

    async def get_tools(self, ctx: Any) -> dict[str, Any]:
        config = get_config()
        tools = await self.wrapped.get_tools(ctx)
        current = get_current_context()
        active_tags = current.active_tags if current else set()

        filtered: dict[str, Any] = {}
        for name, tool in tools.items():
            rules = self._rules_by_name.get(name)
            if rules is None or rules.is_empty():
                if config.require_annotations:
                    raise UserError(f"Tool '{name}' lacks rgltr annotations")
                filtered[name] = tool
                continue
            visible = is_visible(rules, active_tags, config.boundaries)
            record_visibility_event(current, name, rules, active_tags, visible, tool_kind="function")
            if visible:
                filtered[name] = tool
        if current is not None:
            current.audit(
                AuditEvent(
                    event_type="tools_presented",
                    tool_kind="function",
                    tool_name=None,
                    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:
        config = get_config()
        current = get_current_context()
        pre_tags = set(current.active_tags) if current else set()
        rules = self._rules_by_name.get(name)
        if rules is None:
            return await self.wrapped.call_tool(name, tool_args, ctx, tool)
        tool_call_id = getattr(ctx, "tool_call_id", None)

        if current is not None:
            current.audit(
                AuditEvent(
                    event_type="tool_call_requested",
                    tool_kind="function",
                    tool_name=name,
                    tool_call_id=tool_call_id,
                    agent_id=current.agent_id,
                    run_id=current.run_id,
                    active_tags=tags_to_strings(pre_tags),
                )
            )

        visible = is_visible(rules, pre_tags, config.boundaries)
        if not visible:
            record_visibility_event(current, name, rules, pre_tags, visible, tool_kind="function")
            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="function",
                enforced=config.mode != "monitor",
                tool_call_id=tool_call_id,
            )
            if config.mode != "monitor":
                if decision.outcome == "approval":
                    raise ApprovalRequired(metadata={"reason": decision.reason})
                raise ModelRetry(decision.reason)

        would_block = config.mode == "monitor" and decision.outcome == "deny"
        would_require_approval = 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="function",
                    tool_name=name,
                    tool_call_id=tool_call_id,
                    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 config.mode == "monitor" else None,
                    would_require_approval=would_require_approval if config.mode == "monitor" else None,
                )
            )

        return result
