"""
Burrow SDK adapter for AutoGen 0.4+ (Microsoft).

Uses AutoGen's ``DefaultInterventionHandler`` to intercept ``FunctionCall``
and ``FunctionExecutionResult`` messages flowing through the agent runtime,
scanning them with Burrow for prompt injection detection.

Usage:
    from burrow import BurrowGuard
    from burrow.integrations.autogen import create_burrow_handler

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

    from autogen_core import SingleThreadedAgentRuntime

    runtime = SingleThreadedAgentRuntime(
        intervention_handlers=[create_burrow_handler(guard)]
    )
"""

from __future__ import annotations

import json
import logging
from typing import Any

__all__ = [
    "BurrowInterventionHandler",  # noqa: F822 — defined inside create_burrow_handler, hoisted via globals()
    "create_burrow_handler",
]

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


def create_burrow_handler(
    guard: Any,
    block_on_warn: bool = False,
) -> Any:
    """Create an AutoGen 0.4+ intervention handler that scans with Burrow.

    Intercepts ``FunctionCall`` messages (tool invocations) via ``on_send``
    and ``FunctionExecutionResult`` messages (tool outputs) via ``on_response``,
    scanning their content through Burrow for prompt injection detection.

    On block (or warn, if *block_on_warn* is True), the message is dropped
    via ``DropMessage`` so the runtime silently discards it.

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

    Returns:
        A ``DefaultInterventionHandler`` subclass instance.

    Example::

        from autogen_core import SingleThreadedAgentRuntime
        from burrow import BurrowGuard
        from burrow.integrations.autogen import create_burrow_handler

        guard = BurrowGuard(client_id="...", client_secret="...")
        runtime = SingleThreadedAgentRuntime(
            intervention_handlers=[create_burrow_handler(guard)]
        )
    """
    try:
        from autogen_core import (
            AgentId,
            DefaultInterventionHandler,
            DropMessage,
            FunctionCall,
            MessageContext,
        )
        from autogen_core.models import FunctionExecutionResult
    except ImportError as exc:
        raise ImportError(
            "autogen-core is required for the AutoGen adapter. "
            "Install it with: pip install burrow-sdk[autogen]"
        ) from exc

    class BurrowInterventionHandler(DefaultInterventionHandler):
        """AutoGen intervention handler that scans messages with Burrow."""

        def __init__(self, guard: Any, block_on_warn: bool = False) -> None:
            self._guard = guard
            self._block_on_warn = block_on_warn

        async def on_send(
            self,
            message: Any,
            *,
            message_context: MessageContext,
            recipient: AgentId,
        ) -> Any | type[DropMessage]:
            """Intercept outgoing FunctionCall messages and scan tool args."""
            if not isinstance(message, FunctionCall):
                return message

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

            agent_label = f"autogen:{recipient.type}" if recipient else "autogen"

            result = await self._guard.scan_async(
                text,
                content_type="tool_call",
                agent=agent_label,
                tool_name=message.name,
            )

            if result.is_blocked or (self._block_on_warn and result.is_warning):
                logger.warning(
                    "Burrow blocked tool call %s for agent '%s': %s (%.0f%% confidence)",
                    message.name,
                    agent_label,
                    result.category,
                    result.confidence * 100,
                )
                return DropMessage

            return message

        async def on_response(
            self,
            message: Any,
            *,
            sender: AgentId,
            recipient: AgentId | None,
        ) -> Any | type[DropMessage]:
            """Intercept FunctionExecutionResult responses and scan tool output."""
            if not isinstance(message, FunctionExecutionResult):
                return message

            text = message.content if isinstance(message.content, str) else str(message.content)
            if not text.strip():
                return message

            agent_label = f"autogen:{sender.type}" if sender else "autogen"
            tool_name = getattr(message, "name", None) or "unknown"

            result = await self._guard.scan_async(
                text,
                content_type="tool_response",
                agent=agent_label,
                tool_name=tool_name,
            )

            if result.is_blocked or (self._block_on_warn and result.is_warning):
                logger.warning(
                    "Burrow blocked tool response from '%s' (tool=%s): %s (%.0f%% confidence)",
                    agent_label,
                    tool_name,
                    result.category,
                    result.confidence * 100,
                )
                return DropMessage

            return message

    # Hoist the class so users can import it directly if needed.
    globals()["BurrowInterventionHandler"] = BurrowInterventionHandler

    return BurrowInterventionHandler(guard, block_on_warn)


def _extract_text_from_function_call(message: Any) -> str:
    """Extract scannable text from a FunctionCall's arguments."""
    args = getattr(message, "arguments", "")
    if not args:
        return ""
    # arguments is a JSON string in AutoGen's FunctionCall
    if isinstance(args, str):
        try:
            parsed = json.loads(args)
            if isinstance(parsed, dict):
                for key in ("command", "content", "query", "text", "url", "file_path", "pattern"):
                    if key in parsed:
                        return str(parsed[key])
                return args  # scan the full JSON string
            return str(parsed)
        except (json.JSONDecodeError, TypeError):
            return args
    return str(args)
