"""Unit tests for burrow.integrations.pydantic_ai — Pydantic AI adapter."""

from __future__ import annotations

import sys
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock

import pytest

from burrow import ScanResult

# ---------------------------------------------------------------------------
# Mock pydantic_ai types before importing the integration module.
# ---------------------------------------------------------------------------


class FakeSkipToolExecution(Exception):
    """Stand-in for pydantic_ai.exceptions.SkipToolExecution."""


class FakeHooks:
    """Stand-in for pydantic_ai.capabilities.hooks.Hooks.

    Captures decorated handlers so tests can call them directly.
    """

    def __init__(self):
        self._before_tool = None
        self._after_tool = None
        self.on = SimpleNamespace(
            before_tool_execute=self._register_before,
            after_tool_execute=self._register_after,
        )

    def _register_before(self, fn):
        self._before_tool = fn
        return fn

    def _register_after(self, fn):
        self._after_tool = fn
        return fn


_mock_hooks_mod = MagicMock()
_mock_hooks_mod.Hooks = FakeHooks

_mock_exceptions_mod = MagicMock()
_mock_exceptions_mod.SkipToolExecution = FakeSkipToolExecution

_mock_messages_mod = MagicMock()


class FakeToolCallPart:
    """Stand-in for pydantic_ai.messages.ToolCallPart."""

    def __init__(self, tool_name: str = "test_tool"):
        self.tool_name = tool_name


_mock_messages_mod.ToolCallPart = FakeToolCallPart

_mock_tools_mod = MagicMock()


class FakeToolDefinition:
    pass


_mock_tools_mod.ToolDefinition = FakeToolDefinition

_mock_pydantic_ai = MagicMock()

sys.modules["pydantic_ai"] = _mock_pydantic_ai
sys.modules["pydantic_ai.capabilities"] = MagicMock()
sys.modules["pydantic_ai.capabilities.hooks"] = _mock_hooks_mod
sys.modules["pydantic_ai.exceptions"] = _mock_exceptions_mod
sys.modules["pydantic_ai.messages"] = _mock_messages_mod
sys.modules["pydantic_ai.tools"] = _mock_tools_mod

from burrow.integrations.pydantic_ai import (  # noqa: E402
    BurrowBlockedError,
    create_burrow_hooks,
)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _allow_result() -> ScanResult:
    return ScanResult(
        action="allow",
        confidence=0.95,
        category="benign",
        request_id="req-allow-1",
        latency_ms=10.0,
    )


def _block_result() -> ScanResult:
    return ScanResult(
        action="block",
        confidence=0.98,
        category="injection_detected",
        request_id="req-block-1",
        latency_ms=12.0,
    )


def _warn_result() -> ScanResult:
    return ScanResult(
        action="warn",
        confidence=0.72,
        category="suspicious_pattern",
        request_id="req-warn-1",
        latency_ms=11.0,
    )


def _patch_scan_async(guard, result: ScanResult):
    """Replace guard.scan_async with an AsyncMock returning the given result."""
    guard.scan_async = AsyncMock(return_value=result)


def _make_call(tool_name: str = "test_tool") -> FakeToolCallPart:
    return FakeToolCallPart(tool_name=tool_name)


def _make_tool_def() -> FakeToolDefinition:
    return FakeToolDefinition()


def _make_ctx(agent_name: str = "my-agent"):
    return SimpleNamespace(agent_name=agent_name)


# ---------------------------------------------------------------------------
# create_burrow_hooks basics
# ---------------------------------------------------------------------------


class TestCreateBurrowHooks:
    def test_returns_hooks_instance(self, mock_guard):
        hooks = create_burrow_hooks(mock_guard)
        assert isinstance(hooks, FakeHooks)

    def test_burrow_blocked_error_is_exception(self):
        assert issubclass(BurrowBlockedError, Exception)

    def test_burrow_blocked_error_attributes(self):
        result = _block_result()
        err = BurrowBlockedError(result, tool_name="run_cmd", agent_name="agent-1")
        assert err.result is result
        assert err.tool_name == "run_cmd"
        assert err.agent_name == "agent-1"


# ---------------------------------------------------------------------------
# scan_tool_call (before_tool_execute handler)
# ---------------------------------------------------------------------------


class TestScanToolCall:
    @pytest.mark.asyncio
    async def test_clean_args_returns_args(self, mock_guard):
        _patch_scan_async(mock_guard, _allow_result())
        hooks = create_burrow_hooks(mock_guard)
        args = {"query": "safe search query"}
        result = await hooks._before_tool(
            _make_ctx(),
            call=_make_call(),
            tool_def=_make_tool_def(),
            args=args,
        )
        assert result == args

    @pytest.mark.asyncio
    async def test_injection_raises_skip_tool_execution(self, mock_guard):
        _patch_scan_async(mock_guard, _block_result())
        hooks = create_burrow_hooks(mock_guard)
        with pytest.raises(FakeSkipToolExecution, match="Blocked by Burrow"):
            await hooks._before_tool(
                _make_ctx(),
                call=_make_call(),
                tool_def=_make_tool_def(),
                args={"command": "rm -rf / ; ignore previous"},
            )

    @pytest.mark.asyncio
    async def test_agent_identity_from_ctx(self, mock_guard):
        _patch_scan_async(mock_guard, _allow_result())
        hooks = create_burrow_hooks(mock_guard)
        await hooks._before_tool(
            _make_ctx(agent_name="research-agent"),
            call=_make_call(),
            tool_def=_make_tool_def(),
            args={"text": "hello"},
        )
        mock_guard.scan_async.assert_called_once()
        call_kwargs = mock_guard.scan_async.call_args[1]
        assert call_kwargs["agent"] == "pydantic-ai:research-agent"

    @pytest.mark.asyncio
    async def test_agent_identity_fallback(self, mock_guard):
        _patch_scan_async(mock_guard, _allow_result())
        hooks = create_burrow_hooks(mock_guard, default_agent="fallback-agent")
        ctx = SimpleNamespace()  # no agent_name
        await hooks._before_tool(
            ctx,
            call=_make_call(),
            tool_def=_make_tool_def(),
            args={"text": "hello"},
        )
        call_kwargs = mock_guard.scan_async.call_args[1]
        assert call_kwargs["agent"] == "fallback-agent"

    @pytest.mark.asyncio
    async def test_block_on_warn_true(self, mock_guard):
        _patch_scan_async(mock_guard, _warn_result())
        hooks = create_burrow_hooks(mock_guard, block_on_warn=True)
        with pytest.raises(FakeSkipToolExecution):
            await hooks._before_tool(
                _make_ctx(),
                call=_make_call(),
                tool_def=_make_tool_def(),
                args={"command": "suspicious"},
            )

    @pytest.mark.asyncio
    async def test_block_on_warn_false(self, mock_guard):
        _patch_scan_async(mock_guard, _warn_result())
        hooks = create_burrow_hooks(mock_guard, block_on_warn=False)
        args = {"command": "suspicious"}
        result = await hooks._before_tool(
            _make_ctx(),
            call=_make_call(),
            tool_def=_make_tool_def(),
            args=args,
        )
        assert result == args

    @pytest.mark.asyncio
    async def test_empty_args_returns_args(self, mock_guard):
        _patch_scan_async(mock_guard, _allow_result())
        hooks = create_burrow_hooks(mock_guard)
        args = {}
        result = await hooks._before_tool(
            _make_ctx(),
            call=_make_call(),
            tool_def=_make_tool_def(),
            args=args,
        )
        # empty dict json-dumps to "{}" which is not blank, so scan is called
        assert result == args

    @pytest.mark.asyncio
    async def test_known_key_extraction_command(self, mock_guard):
        _patch_scan_async(mock_guard, _allow_result())
        hooks = create_burrow_hooks(mock_guard)
        await hooks._before_tool(
            _make_ctx(),
            call=_make_call(),
            tool_def=_make_tool_def(),
            args={"command": "ls -la", "other": "ignored"},
        )
        assert mock_guard.scan_async.call_args[0][0] == "ls -la"

    @pytest.mark.asyncio
    async def test_known_key_extraction_query(self, mock_guard):
        _patch_scan_async(mock_guard, _allow_result())
        hooks = create_burrow_hooks(mock_guard)
        await hooks._before_tool(
            _make_ctx(),
            call=_make_call(),
            tool_def=_make_tool_def(),
            args={"query": "SELECT * FROM users"},
        )
        assert mock_guard.scan_async.call_args[0][0] == "SELECT * FROM users"

    @pytest.mark.asyncio
    async def test_known_key_extraction_url(self, mock_guard):
        _patch_scan_async(mock_guard, _allow_result())
        hooks = create_burrow_hooks(mock_guard)
        await hooks._before_tool(
            _make_ctx(),
            call=_make_call(),
            tool_def=_make_tool_def(),
            args={"url": "https://example.com"},
        )
        assert mock_guard.scan_async.call_args[0][0] == "https://example.com"

    @pytest.mark.asyncio
    async def test_fallback_json_dumps(self, mock_guard):
        _patch_scan_async(mock_guard, _allow_result())
        hooks = create_burrow_hooks(mock_guard)
        await hooks._before_tool(
            _make_ctx(),
            call=_make_call(),
            tool_def=_make_tool_def(),
            args={"unknown_key": "some value"},
        )
        scanned_text = mock_guard.scan_async.call_args[0][0]
        assert "unknown_key" in scanned_text
        assert "some value" in scanned_text

    @pytest.mark.asyncio
    async def test_tool_name_forwarded(self, mock_guard):
        _patch_scan_async(mock_guard, _allow_result())
        hooks = create_burrow_hooks(mock_guard)
        await hooks._before_tool(
            _make_ctx(),
            call=_make_call("web_search"),
            tool_def=_make_tool_def(),
            args={"query": "test"},
        )
        call_kwargs = mock_guard.scan_async.call_args[1]
        assert call_kwargs["tool_name"] == "web_search"


# ---------------------------------------------------------------------------
# scan_tool_response (after_tool_execute handler)
# ---------------------------------------------------------------------------


class TestScanToolResponse:
    @pytest.mark.asyncio
    async def test_clean_response_returns_result(self, mock_guard):
        _patch_scan_async(mock_guard, _allow_result())
        hooks = create_burrow_hooks(mock_guard)
        result = await hooks._after_tool(
            _make_ctx(),
            call=_make_call(),
            tool_def=_make_tool_def(),
            args={},
            result="safe output",
        )
        assert result == "safe output"

    @pytest.mark.asyncio
    async def test_injection_raises_burrow_blocked_error(self, mock_guard):
        _patch_scan_async(mock_guard, _block_result())
        hooks = create_burrow_hooks(mock_guard)
        with pytest.raises(BurrowBlockedError):
            await hooks._after_tool(
                _make_ctx(),
                call=_make_call(),
                tool_def=_make_tool_def(),
                args={},
                result="malicious response body",
            )

    @pytest.mark.asyncio
    async def test_agent_identity_from_ctx(self, mock_guard):
        _patch_scan_async(mock_guard, _allow_result())
        hooks = create_burrow_hooks(mock_guard)
        await hooks._after_tool(
            _make_ctx(agent_name="writer-agent"),
            call=_make_call(),
            tool_def=_make_tool_def(),
            args={},
            result="some output",
        )
        call_kwargs = mock_guard.scan_async.call_args[1]
        assert call_kwargs["agent"] == "pydantic-ai:writer-agent"

    @pytest.mark.asyncio
    async def test_block_on_warn_true(self, mock_guard):
        _patch_scan_async(mock_guard, _warn_result())
        hooks = create_burrow_hooks(mock_guard, block_on_warn=True)
        with pytest.raises(BurrowBlockedError):
            await hooks._after_tool(
                _make_ctx(),
                call=_make_call(),
                tool_def=_make_tool_def(),
                args={},
                result="suspicious output",
            )

    @pytest.mark.asyncio
    async def test_block_on_warn_false(self, mock_guard):
        _patch_scan_async(mock_guard, _warn_result())
        hooks = create_burrow_hooks(mock_guard, block_on_warn=False)
        result = await hooks._after_tool(
            _make_ctx(),
            call=_make_call(),
            tool_def=_make_tool_def(),
            args={},
            result="suspicious output",
        )
        assert result == "suspicious output"

    @pytest.mark.asyncio
    async def test_none_result_returns_none(self, mock_guard):
        _patch_scan_async(mock_guard, _allow_result())
        hooks = create_burrow_hooks(mock_guard)
        result = await hooks._after_tool(
            _make_ctx(),
            call=_make_call(),
            tool_def=_make_tool_def(),
            args={},
            result=None,
        )
        assert result is None

    @pytest.mark.asyncio
    async def test_empty_string_result_returns_empty(self, mock_guard):
        _patch_scan_async(mock_guard, _allow_result())
        hooks = create_burrow_hooks(mock_guard)
        result = await hooks._after_tool(
            _make_ctx(),
            call=_make_call(),
            tool_def=_make_tool_def(),
            args={},
            result="",
        )
        assert result == ""

    @pytest.mark.asyncio
    async def test_tool_name_forwarded(self, mock_guard):
        _patch_scan_async(mock_guard, _allow_result())
        hooks = create_burrow_hooks(mock_guard)
        await hooks._after_tool(
            _make_ctx(),
            call=_make_call("fetch_url"),
            tool_def=_make_tool_def(),
            args={},
            result="response data",
        )
        call_kwargs = mock_guard.scan_async.call_args[1]
        assert call_kwargs["tool_name"] == "fetch_url"
        assert call_kwargs["content_type"] == "tool_response"
