"""
Burrow SDK adapter for Pydantic AI.

Uses Pydantic AI's Hooks capability (``before_tool_execute`` /
``after_tool_execute``) to scan tool calls and tool responses through
Burrow for prompt injection detection.

Usage:
    from burrow import BurrowGuard
    from burrow.integrations.pydantic_ai import create_burrow_hooks

    guard = BurrowGuard(client_id="...", client_secret="...")
    hooks = create_burrow_hooks(guard)

    agent = Agent("openai:gpt-4o", capabilities=[hooks])
"""

from __future__ import annotations

import json
import logging
from typing import Any

__all__ = ["BurrowBlockedError", "create_burrow_hooks"]

logger = logging.getLogger("burrow.integrations.pydantic_ai")


class BurrowBlockedError(Exception):
    """Raised when Burrow blocks a Pydantic AI tool call or response."""

    def __init__(self, result: Any, tool_name: str = "", agent_name: str = ""):
        self.result = result
        self.tool_name = tool_name
        self.agent_name = agent_name
        super().__init__(
            f"Burrow blocked tool '{tool_name}' for agent '{agent_name}': "
            f"{result.category} ({result.confidence:.0%} confidence)"
        )


def create_burrow_hooks(
    guard: Any,
    block_on_warn: bool = False,
    default_agent: str = "pydantic-ai",
) -> Any:
    """Create Pydantic AI Hooks that scan tool calls and responses via Burrow.

    Returns a ``Hooks`` capability instance with ``before_tool_execute`` and
    ``after_tool_execute`` handlers. Pass it to the agent via the
    ``capabilities`` parameter.

    The ``before_tool_execute`` handler scans tool call arguments
    (``content_type="tool_call"``). If Burrow blocks the call, execution is
    skipped via ``SkipToolExecution`` and a blocked-message is returned to the
    model instead.

    The ``after_tool_execute`` handler scans tool return values
    (``content_type="tool_response"``). If Burrow blocks the response, a
    ``BurrowBlockedError`` is raised.

    Args:
        guard: BurrowGuard instance.
        block_on_warn: If True, also block on "warn" verdicts.
        default_agent: Fallback agent label when agent name is unavailable.

    Returns:
        A ``pydantic_ai.capabilities.hooks.Hooks`` instance.
    """
    try:
        from pydantic_ai.capabilities.hooks import Hooks
        from pydantic_ai.messages import ToolCallPart
        from pydantic_ai.tools import ToolDefinition

        try:
            from pydantic_ai.exceptions import SkipToolExecution
        except ImportError:
            from pydantic_ai.capabilities.hooks import SkipToolExecution
    except ImportError as exc:
        raise ImportError(
            "pydantic-ai is required for the Pydantic AI adapter. "
            "Install it with: pip install burrow-sdk[pydantic-ai]"
        ) from exc

    hooks = Hooks()

    def _resolve_agent(ctx: Any) -> str:
        """Extract agent name from RunContext, if available."""
        agent_name = getattr(ctx, "agent_name", None)
        if not agent_name:
            model = getattr(ctx, "model", None)
            agent_name = getattr(model, "name", None)
        if agent_name:
            return f"pydantic-ai:{agent_name}"
        return default_agent

    def _extract_text(args: dict[str, Any]) -> str:
        """Extract scannable text from tool arguments."""
        for key in ("command", "content", "query", "text", "url", "file_path", "pattern"):
            if key in args:
                return str(args[key])
        return json.dumps(args)

    @hooks.on.before_tool_execute
    async def scan_tool_call(
        ctx: Any,
        *,
        call: ToolCallPart,
        tool_def: ToolDefinition,
        args: dict[str, Any],
    ) -> dict[str, Any]:
        text = _extract_text(args)
        if not text.strip():
            return args

        agent_label = _resolve_agent(ctx)

        result = await guard.scan_async(
            text,
            content_type="tool_call",
            agent=agent_label,
            tool_name=call.tool_name,
            tool_args=args,
        )

        if result.is_blocked or (block_on_warn and result.is_warning):
            logger.warning(
                "Burrow blocked tool call %s for agent '%s': %s (%.0f%% confidence)",
                call.tool_name,
                agent_label,
                result.category,
                result.confidence * 100,
            )
            raise SkipToolExecution(
                f"Blocked by Burrow: {result.category} "
                f"({result.confidence:.0%} confidence). "
                f"DO NOT retry this tool call."
            )

        return args

    @hooks.on.after_tool_execute
    async def scan_tool_response(
        ctx: Any,
        *,
        call: ToolCallPart,
        tool_def: ToolDefinition,
        args: dict[str, Any],
        result: Any,
    ) -> Any:
        text = str(result) if result is not None else ""
        if not text.strip():
            return result

        agent_label = _resolve_agent(ctx)

        scan_result = await guard.scan_async(
            text,
            content_type="tool_response",
            agent=agent_label,
            tool_name=call.tool_name,
        )

        if scan_result.is_blocked or (block_on_warn and scan_result.is_warning):
            logger.warning(
                "Burrow blocked tool response from %s for agent '%s': %s (%.0f%% confidence)",
                call.tool_name,
                agent_label,
                scan_result.category,
                scan_result.confidence * 100,
            )
            raise BurrowBlockedError(scan_result, call.tool_name, agent_label)

        return result

    return hooks
