"""
Burrow adapter for Google Agent Development Kit (ADK).

Provides model-level callbacks (before/after) and tool-level callbacks
(before_tool/after_tool) that scan through Burrow for prompt injection.

Model callbacks see full LLM requests/responses. Tool callbacks see
individual tool calls with ``tool``, ``args``, and ``tool_context``,
enabling more precise scanning.

Usage (model-level)::

    from burrow import BurrowGuard
    from burrow.integrations.adk import create_burrow_callback

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

    agent = Agent(
        model="gemini-2.0-flash",
        before_model_callback=callback,
    )

Usage (full coverage)::

    from burrow.integrations.adk import (
        create_burrow_callback,
        create_burrow_after_callback,
        create_burrow_tool_callback,
        create_burrow_after_tool_callback,
    )

    agent = Agent(
        model="gemini-2.0-flash",
        before_model_callback=create_burrow_callback(guard),
        after_model_callback=create_burrow_after_callback(guard),
        before_tool_callback=create_burrow_tool_callback(guard),
        after_tool_callback=create_burrow_after_tool_callback(guard),
    )
"""

from __future__ import annotations

import logging
from typing import Any

from burrow import BurrowGuard

__all__ = [
    "create_burrow_callback",
    "create_burrow_after_callback",
    "create_burrow_tool_callback",
    "create_burrow_after_tool_callback",
]

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


def create_burrow_callback(
    guard: BurrowGuard,
    block_on_warn: bool = False,
):
    """
    Create a Google ADK before_model_callback with per-agent identity.

    Each ``LlmAgent`` has its own callback slots, so the callback
    reads the agent name from ``callback_context.agent_name`` to produce
    agent identifiers like ``adk:research-agent``.

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

    Returns:
        A callback function compatible with ADK's before_model_callback.
    """
    try:
        from google.adk.models import LlmResponse
        from google.genai.types import Content, Part
    except ImportError as exc:
        raise ImportError(
            "google-adk is required for the ADK adapter. "
            "Install it with: pip install burrow-sdk[adk]"
        ) from exc

    def callback(
        callback_context: Any,
        llm_request: Any,
    ) -> LlmResponse | None:
        if not llm_request.contents:
            return None

        last_content = llm_request.contents[-1]
        if not hasattr(last_content, "parts") or not last_content.parts:
            return None

        text = ""
        for part in last_content.parts:
            if hasattr(part, "text") and part.text:
                text += part.text + " "

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

        ctx_agent_name = getattr(callback_context, "agent_name", "")
        agent_label = f"adk:{ctx_agent_name}" if ctx_agent_name else "adk"

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

        if result.is_blocked or (block_on_warn and result.is_warning):
            return LlmResponse(
                content=Content(
                    role="model",
                    parts=[
                        Part(
                            text=(
                                f"[Blocked by Burrow: {result.category} "
                                f"({result.confidence:.0%} confidence). "
                                f"This message was flagged as a potential "
                                f"prompt injection and was not processed.]"
                            )
                        )
                    ],
                )
            )

        return None

    return callback


def create_burrow_after_callback(
    guard: BurrowGuard,
    block_on_warn: bool = False,
):
    """
    Create a Google ADK after_model_callback with per-agent identity.

    Reads the agent name from ``callback_context.agent_name``.

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

    Returns:
        A callback function compatible with ADK's after_model_callback.
    """
    try:
        from google.adk.models import LlmResponse
        from google.genai.types import Content, Part
    except ImportError as exc:
        raise ImportError(
            "google-adk is required for the ADK adapter. "
            "Install it with: pip install burrow-sdk[adk]"
        ) from exc

    def callback(
        callback_context: Any,
        llm_response: LlmResponse,
    ) -> LlmResponse | None:
        if not llm_response.content or not llm_response.content.parts:
            return None

        text = ""
        for part in llm_response.content.parts:
            if hasattr(part, "text") and part.text:
                text += part.text + " "

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

        ctx_agent_name = getattr(callback_context, "agent_name", "")
        agent_label = f"adk:{ctx_agent_name}" if ctx_agent_name else "adk"

        result = guard.scan(
            text,
            content_type="tool_response",
            agent=agent_label,
        )

        if result.is_blocked or (block_on_warn and result.is_warning):
            return LlmResponse(
                content=Content(
                    role="model",
                    parts=[
                        Part(
                            text=(
                                f"[Blocked by Burrow: Model response contained "
                                f"suspected injection content ({result.category}). "
                                f"Response was not delivered.]"
                            )
                        )
                    ],
                )
            )

        return None

    return callback


def create_burrow_tool_callback(
    guard: BurrowGuard,
    block_on_warn: bool = False,
):
    """
    Create a Google ADK before_tool_callback with per-agent identity.

    Scans individual tool calls before execution, with access to the
    tool object, arguments, and tool context. Reads
    ``tool_context.agent_name`` for per-agent identity like
    ``adk:research-agent``.

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

    Returns:
        A callback compatible with ADK's ``before_tool_callback``.
        Returns a dict to replace the tool result on block, or ``None`` to allow.
    """
    import json as _json

    def callback(
        tool: Any,
        args: dict,
        tool_context: Any,
    ) -> dict | None:
        tool_name = getattr(tool, "name", str(tool))

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

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

        agent_name = getattr(tool_context, "agent_name", "") if tool_context else ""
        agent_label = f"adk:{agent_name}" if agent_name else "adk"

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

        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)",
                tool_name,
                agent_label,
                result.category,
                result.confidence * 100,
            )
            return {
                "status": "error",
                "message": (
                    f"[Blocked by Burrow: {result.category} "
                    f"({result.confidence:.0%} confidence). "
                    f"Tool call was not executed.]"
                ),
            }

        return None

    return callback


def create_burrow_after_tool_callback(
    guard: BurrowGuard,
    block_on_warn: bool = False,
):
    """
    Create a Google ADK after_tool_callback with per-agent identity.

    Scans tool responses after execution for indirect injection.
    Reads ``tool_context.agent_name`` for per-agent identity.

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

    Returns:
        A callback compatible with ADK's ``after_tool_callback``.
        Returns a dict to replace the tool result on block, or ``None`` to allow.
    """

    def callback(
        tool: Any,
        args: dict,
        tool_context: Any,
        tool_response: dict,
    ) -> dict | None:
        tool_name = getattr(tool, "name", str(tool))

        text = str(tool_response)[:4096] if tool_response else ""
        text = text.strip()
        if not text:
            return None

        agent_name = getattr(tool_context, "agent_name", "") if tool_context else ""
        agent_label = f"adk:{agent_name}" if agent_name else "adk"

        result = guard.scan(
            text,
            content_type="tool_response",
            agent=agent_label,
            tool_name=tool_name,
        )

        if result.is_blocked or (block_on_warn and result.is_warning):
            logger.warning(
                "Burrow flagged tool response from '%s' for agent '%s': %s (%.0f%% confidence)",
                tool_name,
                agent_label,
                result.category,
                result.confidence * 100,
            )
            return {
                "status": "error",
                "message": (
                    f"[Blocked by Burrow: Tool output from '{tool_name}' was flagged: "
                    f"{result.category} ({result.confidence:.0%} confidence). "
                    f"Treat with caution.]"
                ),
            }

        return None

    return callback
