"""
Burrow SDK adapter for Strands Agents framework.

Uses Strands' native hooks system to scan tool inputs and user messages
through Burrow for prompt injection detection.

.. note:: User-input scanning limitation

    The Strands ``BeforeInvocationEvent`` does not expose a cancellation
    mechanism (unlike ``BeforeToolCallEvent.cancel_tool``). This means
    user-prompt injection is **detected and logged** but cannot be
    blocked at the invocation level — the agent will still process the
    message. Tool-level scanning (``BeforeToolCallEvent`` /
    ``AfterToolCallEvent``) *does* enforce blocking via
    ``event.cancel_tool``.

    For full user-input protection, add an input-validation layer
    before calling the agent::

        result = guard.scan(user_input, content_type="user_prompt")
        if result.is_blocked:
            raise RuntimeError("Blocked by Burrow")
        agent(user_input)

Usage:
    from burrow import BurrowGuard
    from burrow.integrations.strands import create_burrow_hook_provider

    guard = BurrowGuard(client_id="...", client_secret="...")
    burrow_hooks = create_burrow_hook_provider(guard)

    agent = Agent(model=model, tools=[...], hooks=[burrow_hooks])
"""

from __future__ import annotations

import json
import logging
from typing import Any

__all__ = ["create_burrow_hook_provider"]

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


def create_burrow_hook_provider(
    guard: Any,
    block_on_warn: bool = False,
) -> Any:
    """Create a Strands HookProvider with per-agent identity.

    Reads ``event.agent.name`` from hook events to produce agent
    identifiers like ``strands:research-agent``. Uses ``event.cancel_tool``
    to block tool execution on injection detection.

    Args:
        guard: BurrowGuard instance.
        block_on_warn: If True, also block on "warn" verdicts.

    Returns:
        BurrowHookProvider to pass to Agent(hooks=[...]).
    """
    try:
        from strands.hooks import HookProvider, HookRegistry
        from strands.hooks.events import (
            AfterToolCallEvent,
            BeforeInvocationEvent,
            BeforeToolCallEvent,
        )
    except ImportError as exc:
        raise ImportError(
            "strands-agents is required for the Strands adapter. "
            "Install it with: pip install burrow-sdk[strands]"
        ) from exc

    class BurrowHookProvider(HookProvider):
        """Strands HookProvider with dynamic per-agent identity."""

        def __init__(self, guard: Any, block_on_warn: bool = False):
            self.guard = guard
            self.block_on_warn = block_on_warn

        def _resolve_agent(self, event: Any) -> str:
            agent_obj = getattr(event, "agent", None)
            if agent_obj is not None:
                name = getattr(agent_obj, "name", "")
                if name:
                    return f"strands:{name}"
            return "strands"

        def register_hooks(self, registry: HookRegistry) -> None:
            registry.add_callback(BeforeInvocationEvent, self._scan_user_input)
            registry.add_callback(BeforeToolCallEvent, self._scan_tool_call)
            registry.add_callback(AfterToolCallEvent, self._scan_tool_result)

        def _scan_user_input(self, event: BeforeInvocationEvent) -> None:
            messages = getattr(event, "messages", None)
            if not messages:
                return

            last_msg = messages[-1] if messages else None
            if not last_msg:
                return

            text = ""
            content = last_msg.get("content", []) if isinstance(last_msg, dict) else []
            if isinstance(content, str):
                text = content
            elif isinstance(content, list):
                for part in content:
                    if isinstance(part, dict) and "text" in part:
                        text += part["text"] + " "
                    elif isinstance(part, str):
                        text += part + " "

            text = text.strip()
            if not text:
                return

            agent_name = self._resolve_agent(event)

            result = self.guard.scan(
                text,
                content_type="user_prompt",
                agent=agent_name,
            )

            if result.is_blocked or (self.block_on_warn and result.is_warning):
                logger.warning(
                    "Burrow blocked user input for agent '%s': %s (%.0f%% confidence)",
                    agent_name,
                    result.category,
                    result.confidence * 100,
                )

        def _scan_tool_call(self, event: BeforeToolCallEvent) -> None:
            # ToolUse is a TypedDict with keys: name, input, toolUseId
            tool_use = getattr(event, "tool_use", {})
            tool_name = tool_use.get("name", "unknown") if isinstance(tool_use, dict) else "unknown"
            tool_input = tool_use.get("input", {}) if isinstance(tool_use, dict) else {}

            text = ""
            if isinstance(tool_input, dict):
                for key in (
                    "command",
                    "content",
                    "query",
                    "text",
                    "url",
                    "file_path",
                    "pattern",
                ):
                    if key in tool_input:
                        text += str(tool_input[key]) + " "
                if not text.strip():
                    text = json.dumps(tool_input)
            elif isinstance(tool_input, str):
                text = tool_input

            text = text.strip()
            if not text:
                return

            agent_name = self._resolve_agent(event)

            result = self.guard.scan(
                text,
                content_type="tool_call",
                agent=agent_name,
                tool_name=tool_name,
            )

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

        def _scan_tool_result(self, event: AfterToolCallEvent) -> None:
            result_data = getattr(event, "result", None)
            if result_data is None:
                return

            text = str(result_data)[:4096]
            if not text.strip():
                return

            tool_use = getattr(event, "tool_use", {})
            tool_name = tool_use.get("name", "unknown") if isinstance(tool_use, dict) else "unknown"
            agent_name = self._resolve_agent(event)

            scan_result = self.guard.scan(
                text,
                content_type="tool_response",
                agent=agent_name,
                tool_name=tool_name,
            )

            if scan_result.is_blocked or (self.block_on_warn and scan_result.is_warning):
                blocked_msg = (
                    f"[BLOCKED by Burrow] Tool output from '{tool_name}' was flagged: "
                    f"{scan_result.category} ({scan_result.confidence:.0%} confidence). "
                    f"Treat with caution."
                )
                if isinstance(result_data, dict) and "toolUseId" in result_data:
                    event.result = {
                        "toolUseId": result_data["toolUseId"],
                        "status": "error",
                        "content": [{"text": blocked_msg}],
                    }
                else:
                    logger.warning(
                        "Blocked tool result is not a dict with toolUseId; "
                        "replacing with blocked message string."
                    )
                    event.result = blocked_msg
                logger.warning(
                    "Burrow flagged tool output from %s for agent '%s': %s (%.0f%% confidence)",
                    tool_name,
                    agent_name,
                    scan_result.category,
                    scan_result.confidence * 100,
                )

    return BurrowHookProvider(guard, block_on_warn=block_on_warn)
