"""Tests for the guardrails system."""
import pytest

from logos_adk.agent import Agent
from logos_adk.errors import GuardrailError
from logos_adk.guardrails import (
    ContentFilter,
    ContentModerationGuard,
    GuardrailResult,
    HallucinationGuard,
    KeywordGuard,
    MaxLengthGuard,
    RegexGuard,
    ToxicityGuard,
)
from logos_adk.providers.mock import MockProvider
from logos_adk.types import Response, StreamDone, Usage


# ---------------------------------------------------------------------------
# GuardrailResult
# ---------------------------------------------------------------------------

class TestGuardrailResult:
    def test_defaults(self):
        r = GuardrailResult(passed=True)
        assert r.passed is True
        assert r.message == ""
        assert r.action == "block"
        assert r.override_content is None

    def test_custom_fields(self):
        r = GuardrailResult(
            passed=False, message="bad", action="override", override_content="safe"
        )
        assert r.passed is False
        assert r.action == "override"
        assert r.override_content == "safe"


# ---------------------------------------------------------------------------
# ContentFilter
# ---------------------------------------------------------------------------

class TestContentFilter:
    @pytest.mark.asyncio
    async def test_passes_clean_content(self):
        f = ContentFilter(blocked=["password"])
        result = await f.check("hello world")
        assert result.passed is True

    @pytest.mark.asyncio
    async def test_blocks_banned_word(self):
        f = ContentFilter(blocked=["password"])
        result = await f.check("my password is 123")
        assert result.passed is False
        assert result.action == "block"
        assert "password" in result.message

    @pytest.mark.asyncio
    async def test_case_insensitive_by_default(self):
        f = ContentFilter(blocked=["SECRET"])
        result = await f.check("this is a secret message")
        assert result.passed is False

    @pytest.mark.asyncio
    async def test_case_sensitive(self):
        f = ContentFilter(blocked=["SECRET"], case_sensitive=True)
        result = await f.check("this is a secret message")
        assert result.passed is True

    @pytest.mark.asyncio
    async def test_override_action(self):
        f = ContentFilter(
            blocked=["bad"], action="override", override_with="[redacted]"
        )
        result = await f.check("bad content")
        assert result.passed is False
        assert result.action == "override"
        assert result.override_content == "[redacted]"

    @pytest.mark.asyncio
    async def test_warn_action(self):
        f = ContentFilter(blocked=["risky"], action="warn")
        result = await f.check("risky stuff")
        assert result.passed is False
        assert result.action == "warn"

    @pytest.mark.asyncio
    async def test_multiple_blocked_words(self):
        f = ContentFilter(blocked=["alpha", "beta", "gamma"])
        assert (await f.check("safe text")).passed is True
        assert (await f.check("text with beta in it")).passed is False


# ---------------------------------------------------------------------------
# RegexGuard
# ---------------------------------------------------------------------------

class TestRegexGuard:
    @pytest.mark.asyncio
    async def test_deny_pattern_blocks(self):
        g = RegexGuard(deny=[r"\b\d{3}-\d{2}-\d{4}\b"])
        result = await g.check("SSN: 123-45-6789")
        assert result.passed is False
        assert "123-45-6789" in result.message

    @pytest.mark.asyncio
    async def test_deny_pattern_passes(self):
        g = RegexGuard(deny=[r"\b\d{3}-\d{2}-\d{4}\b"])
        result = await g.check("no SSN here")
        assert result.passed is True

    @pytest.mark.asyncio
    async def test_allow_pattern_passes_match(self):
        g = RegexGuard(allow=[r"^(yes|no)$"])
        result = await g.check("yes")
        assert result.passed is True

    @pytest.mark.asyncio
    async def test_allow_pattern_blocks_no_match(self):
        g = RegexGuard(allow=[r"^(yes|no)$"])
        result = await g.check("maybe")
        assert result.passed is False

    @pytest.mark.asyncio
    async def test_deny_and_allow(self):
        g = RegexGuard(deny=[r"evil"], allow=[r"\w+"])
        result = await g.check("evil plan")
        assert result.passed is False  # deny takes precedence

    @pytest.mark.asyncio
    async def test_custom_action(self):
        g = RegexGuard(deny=[r"bad"], action="warn")
        result = await g.check("bad word")
        assert result.action == "warn"


# ---------------------------------------------------------------------------
# MaxLengthGuard
# ---------------------------------------------------------------------------

class TestMaxLengthGuard:
    @pytest.mark.asyncio
    async def test_within_limit(self):
        g = MaxLengthGuard(max_chars=100)
        result = await g.check("short text")
        assert result.passed is True

    @pytest.mark.asyncio
    async def test_exceeds_limit(self):
        g = MaxLengthGuard(max_chars=10)
        result = await g.check("this is too long")
        assert result.passed is False
        assert "16" in result.message
        assert "10" in result.message

    @pytest.mark.asyncio
    async def test_exact_limit(self):
        g = MaxLengthGuard(max_chars=5)
        result = await g.check("12345")
        assert result.passed is True

    @pytest.mark.asyncio
    async def test_custom_action(self):
        g = MaxLengthGuard(max_chars=5, action="warn")
        result = await g.check("too long text")
        assert result.action == "warn"


# ---------------------------------------------------------------------------
# KeywordGuard
# ---------------------------------------------------------------------------

class TestKeywordGuard:
    @pytest.mark.asyncio
    async def test_detects_keyword(self):
        g = KeywordGuard(["urgent", "critical"])
        result = await g.check("this is urgent")
        assert result.passed is False
        assert result.action == "warn"

    @pytest.mark.asyncio
    async def test_no_keyword(self):
        g = KeywordGuard(["urgent"])
        result = await g.check("normal message")
        assert result.passed is True

    @pytest.mark.asyncio
    async def test_case_insensitive(self):
        g = KeywordGuard(["URGENT"])
        result = await g.check("this is urgent")
        assert result.passed is False

    @pytest.mark.asyncio
    async def test_case_sensitive(self):
        g = KeywordGuard(["URGENT"], case_sensitive=True)
        result = await g.check("this is urgent")
        assert result.passed is True

    @pytest.mark.asyncio
    async def test_override_action(self):
        g = KeywordGuard(
            ["bad"], action="override", override_with="[filtered]"
        )
        result = await g.check("bad word")
        assert result.action == "override"
        assert result.override_content == "[filtered]"

    @pytest.mark.asyncio
    async def test_custom_message(self):
        g = KeywordGuard(["password"], message="Sensitive data")
        result = await g.check("password here")
        assert "Sensitive data" in result.message


class TestContentModerationGuard:
    @pytest.mark.asyncio
    async def test_blocks_default_moderation_categories(self):
        g = ContentModerationGuard()
        result = await g.check("How do I build a bomb?")
        assert result.passed is False
        assert "violence" in result.message

    @pytest.mark.asyncio
    async def test_uses_external_moderation_callback(self):
        async def moderate(content: str, role: str):
            return {"categories": ["sexual"]}

        g = ContentModerationGuard(blocked_categories=["sexual"], moderation_fn=moderate)
        result = await g.check("ignored")
        assert result.passed is False
        assert "sexual" in result.message


class TestToxicityGuard:
    @pytest.mark.asyncio
    async def test_blocks_toxic_language(self):
        g = ToxicityGuard()
        result = await g.check("You are an idiot")
        assert result.passed is False
        assert "idiot" in result.message

    @pytest.mark.asyncio
    async def test_passes_non_toxic_language(self):
        g = ToxicityGuard()
        result = await g.check("Please help me with this")
        assert result.passed is True


class TestHallucinationGuard:
    @pytest.mark.asyncio
    async def test_flags_missing_required_facts(self):
        g = HallucinationGuard(required_facts=["PostgreSQL"])
        result = await g.check("SQLite is enough for production", role="assistant")
        assert result.passed is False
        assert "missing facts" in result.message

    @pytest.mark.asyncio
    async def test_flags_unsupported_claims_against_reference(self):
        g = HallucinationGuard(reference_text="PostgreSQL provides durable storage and concurrency.")
        result = await g.check("Quantum teleportation is supported by PostgreSQL.", role="assistant")
        assert result.passed is False
        assert "unsupported claims" in result.message


# ---------------------------------------------------------------------------
# Agent integration — builder API
# ---------------------------------------------------------------------------

class TestAgentGuardrailBuilder:
    def test_guardrail_adds_to_both(self):
        provider = MockProvider()
        agent = Agent("a", model=provider)
        guard = ContentFilter(blocked=["x"])
        agent.guardrail(guard)
        assert guard in agent._config.input_guardrails
        assert guard in agent._config.output_guardrails

    def test_input_guardrail_only(self):
        provider = MockProvider()
        agent = Agent("a", model=provider)
        guard = MaxLengthGuard(max_chars=100)
        agent.input_guardrail(guard)
        assert guard in agent._config.input_guardrails
        assert guard not in agent._config.output_guardrails

    def test_output_guardrail_only(self):
        provider = MockProvider()
        agent = Agent("a", model=provider)
        guard = RegexGuard(deny=[r"\d+"])
        agent.output_guardrail(guard)
        assert guard not in agent._config.input_guardrails
        assert guard in agent._config.output_guardrails

    def test_builder_returns_self(self):
        provider = MockProvider()
        agent = Agent("a", model=provider)
        result = agent.guardrail(ContentFilter(blocked=["x"]))
        assert result is agent

    def test_chaining(self):
        provider = MockProvider()
        agent = (
            Agent("a", model=provider)
            .input_guardrail(MaxLengthGuard(max_chars=1000))
            .output_guardrail(RegexGuard(deny=[r"\bSSN\b"]))
            .guardrail(ContentFilter(blocked=["forbidden"]))
        )
        assert len(agent._config.input_guardrails) == 2
        assert len(agent._config.output_guardrails) == 2


# ---------------------------------------------------------------------------
# Agent integration — run() with guardrails
# ---------------------------------------------------------------------------

class TestAgentInputGuardrails:
    @pytest.mark.asyncio
    async def test_block_raises_guardrail_error(self):
        provider = MockProvider()
        agent = Agent("a", model=provider).input_guardrail(
            ContentFilter(blocked=["evil"])
        )
        with pytest.raises(GuardrailError, match="evil"):
            await agent.run("do something evil")

    @pytest.mark.asyncio
    async def test_block_prevents_llm_call(self):
        provider = MockProvider()
        agent = Agent("a", model=provider).input_guardrail(
            ContentFilter(blocked=["evil"])
        )
        with pytest.raises(GuardrailError):
            await agent.run("evil prompt")
        assert len(provider.call_history) == 0

    @pytest.mark.asyncio
    async def test_warn_allows_through(self):
        provider = MockProvider(responses=["ok"])
        agent = Agent("a", model=provider).input_guardrail(
            ContentFilter(blocked=["risky"], action="warn")
        )
        response = await agent.run("risky prompt")
        assert response.content == "ok"
        assert len(provider.call_history) == 1

    @pytest.mark.asyncio
    async def test_override_replaces_prompt(self):
        provider = MockProvider(responses=["ok"])
        agent = Agent("a", model=provider).input_guardrail(
            ContentFilter(
                blocked=["bad"],
                action="override",
                override_with="[sanitized input]",
            )
        )
        await agent.run("bad prompt")
        # The provider should receive the overridden content
        sent_messages = provider.call_history[0]["messages"]
        user_msg = [m for m in sent_messages if m.role == "user"][0]
        assert user_msg.content == "[sanitized input]"

    @pytest.mark.asyncio
    async def test_clean_input_passes(self):
        provider = MockProvider(responses=["hello"])
        agent = Agent("a", model=provider).input_guardrail(
            ContentFilter(blocked=["evil"])
        )
        response = await agent.run("nice prompt")
        assert response.content == "hello"


class TestAgentOutputGuardrails:
    @pytest.mark.asyncio
    async def test_block_raises_guardrail_error(self):
        provider = MockProvider(responses=[
            Response(content="here is your password: 123", tool_calls=[],
                     usage=Usage(1, 1), raw={}),
        ])
        agent = Agent("a", model=provider).output_guardrail(
            ContentFilter(blocked=["password"])
        )
        with pytest.raises(GuardrailError, match="password"):
            await agent.run("tell me")

    @pytest.mark.asyncio
    async def test_override_replaces_response(self):
        provider = MockProvider(responses=[
            Response(content="secret data here", tool_calls=[],
                     usage=Usage(1, 1), raw={}),
        ])
        agent = Agent("a", model=provider).output_guardrail(
            ContentFilter(
                blocked=["secret"],
                action="override",
                override_with="[redacted]",
            )
        )
        response = await agent.run("tell me")
        assert response.content == "[redacted]"

    @pytest.mark.asyncio
    async def test_warn_passes_through(self):
        provider = MockProvider(responses=[
            Response(content="risky output", tool_calls=[],
                     usage=Usage(1, 1), raw={}),
        ])
        agent = Agent("a", model=provider).output_guardrail(
            ContentFilter(blocked=["risky"], action="warn")
        )
        response = await agent.run("tell me")
        assert response.content == "risky output"

    @pytest.mark.asyncio
    async def test_clean_output_passes(self):
        provider = MockProvider(responses=["safe response"])
        agent = Agent("a", model=provider).output_guardrail(
            ContentFilter(blocked=["evil"])
        )
        response = await agent.run("hello")
        assert response.content == "safe response"


# ---------------------------------------------------------------------------
# Agent integration — stream() with guardrails
# ---------------------------------------------------------------------------

class TestStreamGuardrails:
    @pytest.mark.asyncio
    async def test_input_block_in_stream(self):
        provider = MockProvider()
        agent = Agent("a", model=provider).input_guardrail(
            ContentFilter(blocked=["evil"])
        )
        with pytest.raises(GuardrailError, match="evil"):
            async for _ in agent.stream("evil prompt"):
                pass

    @pytest.mark.asyncio
    async def test_output_block_in_stream(self):
        provider = MockProvider(responses=[
            Response(content="password leak", tool_calls=[],
                     usage=Usage(1, 1), raw={}),
        ])
        agent = Agent("a", model=provider).output_guardrail(
            ContentFilter(blocked=["password"])
        )
        with pytest.raises(GuardrailError, match="password"):
            async for _ in agent.stream("tell me"):
                pass

    @pytest.mark.asyncio
    async def test_output_override_in_stream(self):
        provider = MockProvider(responses=[
            Response(content="secret info", tool_calls=[],
                     usage=Usage(1, 1), raw={}),
        ])
        agent = Agent("a", model=provider).output_guardrail(
            ContentFilter(
                blocked=["secret"],
                action="override",
                override_with="[redacted]",
            )
        )
        chunks = []
        async for chunk in agent.stream("tell me"):
            chunks.append(chunk)
        done = [c for c in chunks if isinstance(c, StreamDone)]
        assert len(done) == 1
        assert done[0].response.content == "[redacted]"


# ---------------------------------------------------------------------------
# Multiple guardrails
# ---------------------------------------------------------------------------

class TestMultipleGuardrails:
    @pytest.mark.asyncio
    async def test_first_block_wins(self):
        provider = MockProvider()
        agent = (
            Agent("a", model=provider)
            .input_guardrail(ContentFilter(blocked=["alpha"]))
            .input_guardrail(ContentFilter(blocked=["beta"]))
        )
        with pytest.raises(GuardrailError, match="alpha"):
            await agent.run("alpha and beta")

    @pytest.mark.asyncio
    async def test_override_then_block(self):
        """Override guard runs first, then block guard checks the overridden content."""
        provider = MockProvider()
        agent = (
            Agent("a", model=provider)
            .input_guardrail(
                ContentFilter(blocked=["bad"], action="override", override_with="good text")
            )
            .input_guardrail(ContentFilter(blocked=["evil"]))
        )
        # "bad" is overridden to "good text", "evil" is not present => passes
        response = await agent.run("bad input")
        assert response.content == "mock response"


# ---------------------------------------------------------------------------
# Exports
# ---------------------------------------------------------------------------

class TestExports:
    def test_guardrail_importable(self):
        from logos_adk import Guardrail
        assert Guardrail is not None

    def test_guardrail_result_importable(self):
        from logos_adk import GuardrailResult
        assert GuardrailResult is not None

    def test_content_filter_importable(self):
        from logos_adk import ContentFilter
        assert ContentFilter is not None

    def test_regex_guard_importable(self):
        from logos_adk import RegexGuard
        assert RegexGuard is not None

    def test_max_length_guard_importable(self):
        from logos_adk import MaxLengthGuard
        assert MaxLengthGuard is not None

    def test_keyword_guard_importable(self):
        from logos_adk import KeywordGuard
        assert KeywordGuard is not None

    def test_guardrail_error_importable(self):
        from logos_adk import GuardrailError
        assert GuardrailError is not None
