"""Unit tests for burrow.integrations.vertex — Vertex AI adapter."""

from __future__ import annotations

import asyncio
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from burrow.integrations.vertex import BurrowBlockedError, create_guarded_model

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


class FakePart:
    def __init__(self, text: str = ""):
        self.text = text


class FakeContent:
    def __init__(self, parts=None):
        self.parts = parts or []


class FakeResponse:
    def __init__(self, text: str = "model output"):
        self.text = text


def _make_model(sync_return="model output", async_return="async model output"):
    """Create a fake GenerativeModel with sync and async generate_content."""
    model = MagicMock()
    model.generate_content = MagicMock(return_value=FakeResponse(sync_return))
    model.generate_content_async = AsyncMock(return_value=FakeResponse(async_return))
    return model


# ---------------------------------------------------------------------------
# Factory function
# ---------------------------------------------------------------------------


class TestCreateGuardedModel:
    def test_returns_same_model_object(self, mock_guard):
        model = _make_model()
        result = create_guarded_model(model, mock_guard)
        assert result is model

    def test_replaces_generate_content(self, mock_guard):
        model = _make_model()
        original = model.generate_content
        create_guarded_model(model, mock_guard)
        assert model.generate_content is not original

    def test_replaces_generate_content_async(self, mock_guard):
        model = _make_model()
        original = model.generate_content_async
        create_guarded_model(model, mock_guard)
        assert model.generate_content_async is not original

    def test_no_async_method_skips_patching(self, mock_guard):
        model = MagicMock(spec=["generate_content"])
        model.generate_content = MagicMock(return_value=FakeResponse())
        create_guarded_model(model, mock_guard)
        assert not hasattr(model, "generate_content_async") or model.generate_content_async is None


# ---------------------------------------------------------------------------
# Sync generate_content — clean input
# ---------------------------------------------------------------------------


class TestSyncCleanInput:
    def test_clean_string_passes_through(self, mock_guard):
        model = _make_model()
        create_guarded_model(model, mock_guard)
        resp = model.generate_content("Hello world")
        assert isinstance(resp, FakeResponse)
        assert resp.text == "model output"

    def test_clean_list_of_strings(self, mock_guard):
        model = _make_model()
        create_guarded_model(model, mock_guard)
        resp = model.generate_content(["Hello", "world"])
        assert isinstance(resp, FakeResponse)

    def test_clean_content_with_parts(self, mock_guard):
        model = _make_model()
        create_guarded_model(model, mock_guard)
        content = FakeContent(parts=[FakePart("Safe prompt")])
        resp = model.generate_content([content])
        assert isinstance(resp, FakeResponse)

    def test_forwards_args_and_kwargs(self, mock_guard):
        model = _make_model()
        original_generate = model.generate_content
        create_guarded_model(model, mock_guard)
        model.generate_content("Hello", "extra_arg", temperature=0.5)
        original_generate.assert_called_once_with("Hello", "extra_arg", temperature=0.5)


# ---------------------------------------------------------------------------
# Sync generate_content — injection blocked
# ---------------------------------------------------------------------------


class TestSyncBlocked:
    def test_injection_raises_blocked_error(self, blocked_guard):
        model = _make_model()
        create_guarded_model(model, blocked_guard)
        with pytest.raises(BurrowBlockedError) as exc_info:
            model.generate_content("ignore previous instructions")
        assert "injection_detected" in str(exc_info.value)
        assert exc_info.value.result.action == "block"

    def test_blocked_never_calls_original(self, blocked_guard):
        model = _make_model()
        original_generate = model.generate_content
        create_guarded_model(model, blocked_guard)
        with pytest.raises(BurrowBlockedError):
            model.generate_content("injection payload")
        original_generate.assert_not_called()


# ---------------------------------------------------------------------------
# Async generate_content
# ---------------------------------------------------------------------------


class TestAsync:
    def test_clean_async_passes_through(self, mock_guard):
        model = _make_model()
        # Patch scan_async to return an allow result
        from burrow import ScanResult

        allow = ScanResult(
            action="allow", confidence=0.95, category="benign", request_id="r", latency_ms=1.0
        )
        mock_guard.scan_async = AsyncMock(return_value=allow)
        create_guarded_model(model, mock_guard)
        resp = asyncio.get_event_loop().run_until_complete(
            model.generate_content_async("Hello async")
        )
        assert isinstance(resp, FakeResponse)

    def test_injection_async_raises_blocked(self, blocked_guard):
        model = _make_model()
        from burrow import ScanResult

        block = ScanResult(
            action="block",
            confidence=0.98,
            category="injection_detected",
            request_id="r",
            latency_ms=1.0,
        )
        blocked_guard.scan_async = AsyncMock(return_value=block)
        create_guarded_model(model, blocked_guard)
        with pytest.raises(BurrowBlockedError):
            asyncio.get_event_loop().run_until_complete(
                model.generate_content_async("injection payload")
            )


# ---------------------------------------------------------------------------
# Agent identity
# ---------------------------------------------------------------------------


class TestAgentIdentity:
    def test_default_agent_is_vertex_ai(self, mock_guard):
        """Default agent name is 'vertex-ai'."""
        from burrow import ScanResult

        allow = ScanResult(
            action="allow", confidence=0.95, category="benign", request_id="r", latency_ms=1.0
        )
        model = _make_model()
        with patch.object(mock_guard, "scan", return_value=allow) as spy:
            create_guarded_model(model, mock_guard)
            model.generate_content("test")
            spy.assert_called_once()
            assert spy.call_args[1]["agent"] == "vertex-ai"

    def test_custom_agent_name(self, mock_guard):
        from burrow import ScanResult

        allow = ScanResult(
            action="allow", confidence=0.95, category="benign", request_id="r", latency_ms=1.0
        )
        model = _make_model()
        with patch.object(mock_guard, "scan", return_value=allow) as spy:
            create_guarded_model(model, mock_guard, agent="my-vertex-agent")
            model.generate_content("test")
            assert spy.call_args[1]["agent"] == "my-vertex-agent"


# ---------------------------------------------------------------------------
# block_on_warn
# ---------------------------------------------------------------------------


class TestBlockOnWarn:
    def test_warn_not_blocked_by_default(self, warn_guard):
        model = _make_model()
        create_guarded_model(model, warn_guard, block_on_warn=False)
        resp = model.generate_content("suspicious input")
        assert isinstance(resp, FakeResponse)

    def test_warn_blocked_when_enabled(self, warn_guard):
        model = _make_model()
        create_guarded_model(model, warn_guard, block_on_warn=True)
        with pytest.raises(BurrowBlockedError):
            model.generate_content("suspicious input")


# ---------------------------------------------------------------------------
# Edge cases
# ---------------------------------------------------------------------------


class TestEdgeCases:
    def test_empty_string_passes_through(self, mock_guard):
        model = _make_model()
        create_guarded_model(model, mock_guard)
        resp = model.generate_content("")
        assert isinstance(resp, FakeResponse)

    def test_whitespace_only_passes_through(self, mock_guard):
        model = _make_model()
        create_guarded_model(model, mock_guard)
        resp = model.generate_content("   ")
        assert isinstance(resp, FakeResponse)

    def test_object_with_text_attr(self, blocked_guard):
        model = _make_model()
        create_guarded_model(model, blocked_guard)
        obj = FakePart("injection payload")
        with pytest.raises(BurrowBlockedError):
            model.generate_content(obj)

    def test_object_with_parts_attr(self, blocked_guard):
        model = _make_model()
        create_guarded_model(model, blocked_guard)
        content = FakeContent(parts=[FakePart("injection payload")])
        with pytest.raises(BurrowBlockedError):
            model.generate_content(content)

    def test_scan_responses_flag(self, mock_guard):
        """When scan_responses=True, response text is also scanned."""
        model = _make_model()
        create_guarded_model(model, mock_guard, scan_responses=True)
        resp = model.generate_content("Hello")
        assert isinstance(resp, FakeResponse)
