"""PolicyGatedTools — structural policy enforcement at the tool layer.

Wraps a ToolProvider to filter and gate tool access based on graph policies.
Instead of hoping the LLM follows prompt instructions, this controls what
tools the LLM *can see* and *can execute*.

- ``definitions()`` returns only tools allowed by active policies.
- ``execute()`` blocks calls to tools not permitted by policy.
- When no policies exist yet (cold start), all tools are available.
"""

from __future__ import annotations

import re
from typing import TYPE_CHECKING, Any

import structlog

from koan.tools.types import ToolResult

if TYPE_CHECKING:
    from koan.tools.types import ToolDefinition, ToolProvider

log = structlog.get_logger()


class PolicyGatedTools:
    """ToolProvider wrapper that enforces graph policies structurally.

    Usage::

        gated = PolicyGatedTools(registry, allowed_tools=["lookup_account"])
        agent = ReactAgent("bot", client=llm, tools=gated)
        # LLM only sees lookup_account — can't even attempt other tools

    When ``allowed_tools`` is None or empty, all tools pass through
    (permissive mode for cold start before any policies exist).
    """

    def __init__(
        self,
        inner: ToolProvider,
        *,
        allowed_tools: list[str] | None = None,
    ) -> None:
        self._inner = inner
        self._allowed: set[str] | None = (
            set(allowed_tools) if allowed_tools is not None else None
        )

    def definitions(self) -> list[ToolDefinition]:
        """Return only tool definitions permitted by policy.

        If no policy constraints exist (allowed is None), returns all tools.
        """
        all_defs = self._inner.definitions()
        if self._allowed is None:
            return all_defs

        filtered = [d for d in all_defs if d.name in self._allowed]

        if len(filtered) < len(all_defs):
            blocked = {d.name for d in all_defs} - self._allowed
            log.debug(
                "gated_tools.filtered",
                allowed=len(filtered),
                blocked=sorted(blocked),
            )

        return filtered

    async def execute(
        self,
        name: str,
        arguments: dict[str, Any] | str,
        *,
        context: Any | None = None,
    ) -> ToolResult:
        """Execute a tool, blocking calls not permitted by policy."""
        if self._allowed is not None and name not in self._allowed:
            log.warning("gated_tools.blocked", tool=name)
            return ToolResult(
                tool_name=name,
                error=(
                    f"Tool '{name}' is not permitted by current policies. "
                    f"Allowed tools: {sorted(self._allowed)}"
                ),
            )

        return await self._inner.execute(name, arguments, context=context)

    @staticmethod
    def extract_tools_from_policies(
        policies: list[dict[str, Any]],
    ) -> list[str] | None:
        """Parse policy rules to extract allowed tool names.

        Returns a list of tool names if policies constrain tools,
        or None if no tool constraints exist (permissive mode).

        Parses rules like:
          - "For classification decisions, use: lookup_account"
          - "For classification decisions, use: lookup_account then apply_credit"
        """
        tool_names: set[str] = set()

        for policy in policies:
            rules = policy.get("rules", [])
            for rule in rules:
                if not isinstance(rule, str):
                    continue
                # Match "use: tool_a" or "use: tool_a then tool_b"
                match = re.search(r"use:\s*(.+)", rule, re.IGNORECASE)
                if match:
                    tools_part = match.group(1).strip()
                    for t in re.split(r"\s+then\s+", tools_part):
                        cleaned = t.strip().rstrip(".,;")
                        if cleaned and cleaned != "no specific tool":
                            tool_names.add(cleaned)

        return sorted(tool_names) if tool_names else None
