"""Tests for guardrails — input/output validation, tripwire semantics, advisory streaming."""

import json
from types import SimpleNamespace
from typing import Any, AsyncIterator, cast
from unittest.mock import AsyncMock, patch

import pytest

from swarm.agent import Agent
from swarm.chunk import ChunkType, DelimiterChunk
from swarm.context import RunContext
from swarm.guardrails import (
    GuardrailResult,
    InputGuardrail,
    InputGuardrailTripped,
    OutputGuardrail,
    OutputGuardrailContext,
    OutputGuardrailTripped,
    _extract_last_assistant_text,
    run_input_guardrails,
    run_output_guardrails,
)
from swarm.message import SwarmToolMessage
from swarm.providers.protocols import ModelRouterProtocol
from swarm.runners.agent import AgentRunner, AgentSignal
from swarm.runners.dialogue import DialogueRunner
from swarm.tool import ResultMetadata, ToolExecutionResponse
from swarm.types import ContentType

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


async def _pass_guardrail(*args: Any, **kwargs: Any) -> GuardrailResult:
    return GuardrailResult(tripwire_triggered=False)


async def _trip_guardrail(*args: Any, **kwargs: Any) -> GuardrailResult:
    return GuardrailResult(tripwire_triggered=True, output_info="blocked")


def _make_agent(
    *,
    input_guardrails: list[Any] | None = None,
    output_guardrails: list[Any] | None = None,
) -> Agent:
    return Agent(
        name="GuardrailTest",
        model="gpt-4o",
        instructions="test",
        tools=[],
        input_guardrails=input_guardrails or [],
        output_guardrails=output_guardrails or [],
    )


class _FakeContext:
    conversation_id = "conv-1"
    current_context_variables: dict[str, Any] = {}
    workflow_context: dict[str, Any] = {}
    designer_fullname = "Agent"

    def get_history_for_completion(self) -> list[Any]:
        return []


class _FakeStreamManager:
    async def publish_status(self, status: str) -> None:
        pass


class _FakeMessageStore:
    async def persist_messages(self, **kwargs: Any) -> None:
        pass


class _FakeRouter:
    def __init__(self, responses: list[Any]) -> None:
        self._responses = list(responses)

    async def create_completion_async(self, request: Any) -> Any:
        return self._responses.pop(0)

    async def stream_completion(self, request: Any) -> AsyncIterator[Any]:
        _ = request
        raise AssertionError(
            "stream_completion should not be called in AgentRunner tests"
        )
        yield None  # pragma: no cover


def _make_completion(*, content: str = "done") -> Any:
    return SimpleNamespace(
        choices=[
            SimpleNamespace(message=SimpleNamespace(content=content, tool_calls=None))
        ],
        usage=None,
        model="gpt-4o",
    )


def _make_run_context(router: _FakeRouter) -> RunContext[_FakeContext]:
    return RunContext(
        deps=_FakeContext(),
        stream_manager=_FakeStreamManager(),
        message_store=_FakeMessageStore(),
        router=router,
        retry_policy=None,
    )


# ---------------------------------------------------------------------------
# _extract_last_assistant_text
# ---------------------------------------------------------------------------


class TestExtractLastAssistantText:
    def test_extracts_from_dict_history(self) -> None:
        history = [
            {"role": "user", "content": [{"type": "text", "text": "hello"}]},
            {"role": "assistant", "content": [{"type": "text", "text": "world"}]},
            {"role": "tool", "content": [{"type": "text", "text": "result"}]},
        ]
        assert _extract_last_assistant_text(history) == "world"

    def test_extracts_from_object_history(self) -> None:
        msg = SimpleNamespace(
            role="assistant",
            content=[
                {"type": "text", "text": "line1"},
                {"type": "text", "text": "line2"},
            ],
        )
        history = [SimpleNamespace(role="user", content=[]), msg]
        assert _extract_last_assistant_text(history) == "line1\nline2"

    def test_ignores_non_text_parts(self) -> None:
        msg = SimpleNamespace(
            role="assistant",
            content=[
                {"type": "text", "text": "good"},
                {"type": "image", "url": "x.png"},
            ],
        )
        assert _extract_last_assistant_text([msg]) == "good"

    def test_backward_scan_finds_last_assistant(self) -> None:
        history = [
            {"role": "assistant", "content": [{"type": "text", "text": "first"}]},
            {"role": "user", "content": []},
            {"role": "assistant", "content": [{"type": "text", "text": "second"}]},
            {"role": "tool", "content": []},
        ]
        assert _extract_last_assistant_text(history) == "second"

    def test_empty_history_returns_empty_string(self) -> None:
        assert _extract_last_assistant_text([]) == ""

    def test_no_assistant_message_returns_empty_string(self) -> None:
        history = [{"role": "user", "content": []}]
        assert _extract_last_assistant_text(history) == ""


# ---------------------------------------------------------------------------
# run_input_guardrails
# ---------------------------------------------------------------------------


class TestRunInputGuardrails:
    @pytest.mark.asyncio
    async def test_passes_when_no_guardrails(self) -> None:
        agent = _make_agent()
        results = await run_input_guardrails(agent, [], {})
        assert results == []

    @pytest.mark.asyncio
    async def test_passes_when_guardrail_returns_no_trip(self) -> None:
        agent = _make_agent(
            input_guardrails=[InputGuardrail(name="safe", guardrail_fn=_pass_guardrail)]
        )
        results = await run_input_guardrails(agent, [], {})
        assert len(results) == 1
        assert not results[0].tripwire_triggered

    @pytest.mark.asyncio
    async def test_raises_on_trip(self) -> None:
        agent = _make_agent(
            input_guardrails=[
                InputGuardrail(name="blocker", guardrail_fn=_trip_guardrail)
            ]
        )
        with pytest.raises(InputGuardrailTripped) as exc_info:
            await run_input_guardrails(agent, [], {})
        assert exc_info.value.guardrail_name == "blocker"
        assert exc_info.value.result.tripwire_triggered

    @pytest.mark.asyncio
    async def test_sequential_first_trip_halts_remaining(self) -> None:
        call_order: list[str] = []

        async def first_fn(*args: Any) -> GuardrailResult:
            call_order.append("first")
            return GuardrailResult(tripwire_triggered=True)

        async def second_fn(*args: Any) -> GuardrailResult:
            call_order.append("second")
            return GuardrailResult(tripwire_triggered=False)

        agent = _make_agent(
            input_guardrails=[
                InputGuardrail(name="first", guardrail_fn=first_fn),
                InputGuardrail(name="second", guardrail_fn=second_fn),
            ]
        )
        with pytest.raises(InputGuardrailTripped):
            await run_input_guardrails(agent, [], {})
        assert call_order == ["first"]  # second never called

    @pytest.mark.asyncio
    async def test_guardrail_name_set_authoritatively(self) -> None:
        """Runner stamps guardrail_name, not the fn."""

        async def fn_with_wrong_name(*args: Any) -> GuardrailResult:
            return GuardrailResult(guardrail_name="wrong")

        agent = _make_agent(
            input_guardrails=[
                InputGuardrail(name="correct", guardrail_fn=fn_with_wrong_name)
            ]
        )
        results = await run_input_guardrails(agent, [], {})
        assert results[0].guardrail_name == "correct"


# ---------------------------------------------------------------------------
# run_output_guardrails
# ---------------------------------------------------------------------------


class TestRunOutputGuardrails:
    @pytest.mark.asyncio
    async def test_passes_when_no_guardrails(self) -> None:
        agent = _make_agent()
        ctx = OutputGuardrailContext(
            agent_name="test", agent_output="hi", context_variables={}
        )
        results = await run_output_guardrails(agent, ctx)
        assert results == []

    @pytest.mark.asyncio
    async def test_passes_when_guardrail_returns_no_trip(self) -> None:
        agent = _make_agent(
            output_guardrails=[
                OutputGuardrail(name="safe", guardrail_fn=_pass_guardrail)
            ]
        )
        ctx = OutputGuardrailContext(
            agent_name="test", agent_output="hi", context_variables={}
        )
        results = await run_output_guardrails(agent, ctx)
        assert len(results) == 1
        assert not results[0].tripwire_triggered

    @pytest.mark.asyncio
    async def test_trip_halts_remaining(self) -> None:
        call_order: list[str] = []

        async def first_fn(*args: Any) -> GuardrailResult:
            call_order.append("first")
            return GuardrailResult(tripwire_triggered=True, output_info="bad")

        async def second_fn(*args: Any) -> GuardrailResult:
            call_order.append("second")
            return GuardrailResult()

        agent = _make_agent(
            output_guardrails=[
                OutputGuardrail(name="first", guardrail_fn=first_fn),
                OutputGuardrail(name="second", guardrail_fn=second_fn),
            ]
        )
        ctx = OutputGuardrailContext(
            agent_name="test", agent_output="bad", context_variables={}
        )
        results = await run_output_guardrails(agent, ctx)
        assert call_order == ["first"]
        assert len(results) == 1
        assert results[0].tripwire_triggered

    @pytest.mark.asyncio
    async def test_guardrail_name_set_authoritatively(self) -> None:
        async def fn(*args: Any) -> GuardrailResult:
            return GuardrailResult(guardrail_name="wrong")

        agent = _make_agent(
            output_guardrails=[OutputGuardrail(name="correct", guardrail_fn=fn)]
        )
        ctx = OutputGuardrailContext(
            agent_name="test", agent_output="hi", context_variables={}
        )
        results = await run_output_guardrails(agent, ctx)
        assert results[0].guardrail_name == "correct"


# ---------------------------------------------------------------------------
# Batch runner integration — guardrails
# ---------------------------------------------------------------------------


class TestAgentRunnerGuardrails:
    @pytest.mark.asyncio
    async def test_input_guardrail_passes_llm_proceeds(self) -> None:
        agent = _make_agent(
            input_guardrails=[InputGuardrail(name="safe", guardrail_fn=_pass_guardrail)]
        )
        router = _FakeRouter([_make_completion(content="done")])
        run_context = _make_run_context(router)

        runner = AgentRunner(execute_tools_fn=cast(Any, AsyncMock()))
        result = await runner.run_batch(
            agent=agent, run_context=cast(Any, run_context), max_turns=1
        )
        assert result.signal == AgentSignal.COMPLETED

    @pytest.mark.asyncio
    async def test_input_guardrail_trips_before_llm_call(self) -> None:
        agent = _make_agent(
            input_guardrails=[
                InputGuardrail(name="blocker", guardrail_fn=_trip_guardrail)
            ]
        )
        router = _FakeRouter([])  # no completions — should never be reached
        run_context = _make_run_context(router)

        runner = AgentRunner(execute_tools_fn=cast(Any, AsyncMock()))
        with pytest.raises(InputGuardrailTripped):
            await runner.run_batch(
                agent=agent, run_context=cast(Any, run_context), max_turns=1
            )

    @pytest.mark.asyncio
    async def test_output_guardrail_passes_result_returned(self) -> None:
        agent = _make_agent(
            output_guardrails=[
                OutputGuardrail(name="safe", guardrail_fn=_pass_guardrail)
            ]
        )
        router = _FakeRouter([_make_completion(content="done")])
        run_context = _make_run_context(router)

        runner = AgentRunner(execute_tools_fn=cast(Any, AsyncMock()))
        result = await runner.run_batch(
            agent=agent, run_context=cast(Any, run_context), max_turns=1
        )
        assert result.signal == AgentSignal.COMPLETED

    @pytest.mark.asyncio
    async def test_output_guardrail_trips_raises(self) -> None:
        agent = _make_agent(
            output_guardrails=[
                OutputGuardrail(name="blocker", guardrail_fn=_trip_guardrail)
            ]
        )
        router = _FakeRouter([_make_completion(content="done")])
        run_context = _make_run_context(router)

        runner = AgentRunner(execute_tools_fn=cast(Any, AsyncMock()))
        with pytest.raises(OutputGuardrailTripped):
            await runner.run_batch(
                agent=agent, run_context=cast(Any, run_context), max_turns=1
            )

    @pytest.mark.asyncio
    async def test_agent_with_no_guardrails_works(self) -> None:
        agent = _make_agent()
        router = _FakeRouter([_make_completion(content="done")])
        run_context = _make_run_context(router)

        runner = AgentRunner(execute_tools_fn=cast(Any, AsyncMock()))
        result = await runner.run_batch(
            agent=agent, run_context=cast(Any, run_context), max_turns=1
        )
        assert result.signal == AgentSignal.COMPLETED


# ---------------------------------------------------------------------------
# Streaming runner integration — advisory-only output guardrails
# ---------------------------------------------------------------------------


class _FakeStreamDelta:
    """Mimics chunk.choices[0].delta with a .json() method."""

    def __init__(self, content: str) -> None:
        self._content = content

    def json(self) -> str:
        return json.dumps({"role": "assistant", "content": self._content})


class _FakeStreamChunk:
    """Mimics a streaming chunk from router.stream_completion."""

    def __init__(self, content: str) -> None:
        self.choices = [SimpleNamespace(delta=_FakeStreamDelta(content))]
        self.usage = None


class _FakeStreamRouter:
    """Router that yields stream chunks for DialogueRunner."""

    def __init__(self, chunks: list[_FakeStreamChunk]) -> None:
        self._chunks = chunks

    async def create_completion_async(self, request: Any) -> Any:
        _ = request
        raise AssertionError(
            "create_completion_async should not be called in DialogueRunner tests"
        )

    async def stream_completion(self, request: Any) -> AsyncIterator[Any]:
        _ = request
        for chunk in self._chunks:
            yield chunk


class _FakeStreamContext:
    """Context for streaming tests."""

    conversation_id = "conv-1"
    current_context_variables: dict[str, Any] = {}
    workflow_context: dict[str, Any] = {}
    designer_fullname = "Agent"

    def get_history_for_completion(self) -> list[Any]:
        return []


def _make_stream_run_context(
    router: ModelRouterProtocol,
) -> RunContext[_FakeStreamContext]:
    return RunContext(
        deps=_FakeStreamContext(),
        stream_manager=_FakeStreamManager(),
        message_store=_FakeMessageStore(),
        router=router,
        retry_policy=None,
    )


class TestDialogueRunnerGuardrails:
    """Streaming output guardrails are advisory-only: no exception raised,
    completion chunk still emitted, results on ExecutionMetadata."""

    @pytest.mark.asyncio
    async def test_output_guardrail_trip_remains_advisory(self) -> None:
        """A tripping output guardrail must NOT raise — the final
        DelimiterChunk(delim='complete') must still be yielded with the
        guardrail results attached to execution_metadata."""
        agent = _make_agent(
            output_guardrails=[
                OutputGuardrail(name="content-filter", guardrail_fn=_trip_guardrail)
            ]
        )
        router = _FakeStreamRouter([_FakeStreamChunk("hello world")])
        run_context = _make_stream_run_context(router)

        runner = DialogueRunner()
        chunks: list[ChunkType] = []
        async for chunk in runner.run_streamed(
            agent=agent, run_context=cast(Any, run_context), max_turns=1
        ):
            chunks.append(chunk)

        # Final chunk must be DelimiterChunk(delim="complete")
        final = chunks[-1]
        assert isinstance(final, DelimiterChunk)
        assert final.delim == "complete"

        # execution_metadata must carry the tripped guardrail result
        assert final.execution_metadata is not None
        gr_results = final.execution_metadata.output_guardrail_results
        assert len(gr_results) == 1
        assert gr_results[0].tripwire_triggered is True
        assert gr_results[0].guardrail_name == "content-filter"
        assert gr_results[0].output_info == "blocked"

    @pytest.mark.asyncio
    async def test_output_guardrail_pass_returns_results(self) -> None:
        """A passing output guardrail still populates execution_metadata."""
        agent = _make_agent(
            output_guardrails=[
                OutputGuardrail(name="safe-check", guardrail_fn=_pass_guardrail)
            ]
        )
        router = _FakeStreamRouter([_FakeStreamChunk("safe content")])
        run_context = _make_stream_run_context(router)

        runner = DialogueRunner()
        chunks: list[ChunkType] = []
        async for chunk in runner.run_streamed(
            agent=agent, run_context=cast(Any, run_context), max_turns=1
        ):
            chunks.append(chunk)

        final = chunks[-1]
        assert isinstance(final, DelimiterChunk)
        assert final.delim == "complete"
        assert final.execution_metadata is not None
        gr_results = final.execution_metadata.output_guardrail_results
        assert len(gr_results) == 1
        assert gr_results[0].tripwire_triggered is False

    @pytest.mark.asyncio
    async def test_no_output_guardrails_yields_empty_results(self) -> None:
        """Agent without output guardrails yields empty results list."""
        agent = _make_agent()
        router = _FakeStreamRouter([_FakeStreamChunk("done")])
        run_context = _make_stream_run_context(router)

        runner = DialogueRunner()
        chunks: list[ChunkType] = []
        async for chunk in runner.run_streamed(
            agent=agent, run_context=cast(Any, run_context), max_turns=1
        ):
            chunks.append(chunk)

        final = chunks[-1]
        assert isinstance(final, DelimiterChunk)
        assert final.delim == "complete"
        assert final.execution_metadata is not None
        assert final.execution_metadata.output_guardrail_results == []

    @pytest.mark.asyncio
    async def test_input_guardrail_trips_before_streaming(self) -> None:
        """A tripping input guardrail raises before any streaming occurs."""
        agent = _make_agent(
            input_guardrails=[
                InputGuardrail(name="blocker", guardrail_fn=_trip_guardrail)
            ]
        )
        router = _FakeStreamRouter([])  # no chunks — should never be reached
        run_context = _make_stream_run_context(router)

        runner = DialogueRunner()
        with pytest.raises(InputGuardrailTripped):
            async for _ in runner.run_streamed(
                agent=agent, run_context=cast(Any, run_context), max_turns=1
            ):
                pass

    @pytest.mark.asyncio
    async def test_max_turns_zero_does_not_crash(self) -> None:
        """max_turns=0 means the loop body never executes. The runner must
        still yield the final DelimiterChunk without UnboundLocalError."""
        agent = _make_agent(
            output_guardrails=[
                OutputGuardrail(name="safe", guardrail_fn=_pass_guardrail)
            ]
        )
        router = _FakeStreamRouter([])  # no chunks — loop never enters
        run_context = _make_stream_run_context(router)

        runner = DialogueRunner()
        chunks: list[ChunkType] = []
        async for chunk in runner.run_streamed(
            agent=agent, run_context=cast(Any, run_context), max_turns=0
        ):
            chunks.append(chunk)

        final = chunks[-1]
        assert isinstance(final, DelimiterChunk)
        assert final.delim == "complete"

    @pytest.mark.asyncio
    async def test_output_guardrail_uses_producing_agent_not_handoff_target(
        self,
    ) -> None:
        """After a save-tool handoff, output guardrails must run against
        the agent that *produced* the output, not the handoff target."""
        guardrail_received_agent: list[Any] = []

        async def recording_guardrail(agent: Any, *args: Any) -> GuardrailResult:
            guardrail_received_agent.append(agent)
            return GuardrailResult(tripwire_triggered=True, output_info="recorded")

        # Agent A has an output guardrail
        agent_a = Agent(
            name="AgentA",
            model="gpt-4o",
            instructions="test",
            tools=[],
            output_guardrails=[
                OutputGuardrail(name="recorder", guardrail_fn=recording_guardrail)
            ],
        )

        # Agent B has NO output guardrails
        agent_b = Agent(
            name="AgentB",
            model="gpt-4o",
            instructions="test",
            tools=[],
            output_guardrails=[],
        )

        # Stream: first chunk is assistant text, second is a tool call
        text_chunk = _FakeStreamChunk("Agent A speaking")

        tool_delta = {
            "role": "assistant",
            "tool_calls": [
                {
                    "index": 0,
                    "id": "call-1",
                    "type": "function",
                    "function": {"name": "save_data", "arguments": "{}"},
                }
            ],
        }

        class _ToolCallDelta:
            def json(self) -> str:
                return json.dumps(tool_delta)

        tool_chunk = SimpleNamespace(
            choices=[SimpleNamespace(delta=_ToolCallDelta())],
            usage=None,
        )

        router = _FakeStreamRouter(cast(Any, [text_chunk, tool_chunk]))
        run_context = _make_stream_run_context(router)

        # Fake execute_tools: returns a save-tool response with handoff to Agent B
        # and agent_paused=True so the loop exits after the handoff
        tool_msg = SwarmToolMessage(
            id="tool-msg-1",
            tool_call_id="call-1",
            tool_name="save_data",
            content="saved",
            content_type=ContentType.TEXT,
        )

        fake_tool_response = ToolExecutionResponse(
            messages=[tool_msg],
            agent=agent_b,
            context_variables={},
            metadata=ResultMetadata(agent_paused=True),
        )

        async def fake_execute_tools(
            *args: Any, **kwargs: Any
        ) -> ToolExecutionResponse:
            return fake_tool_response

        runner = DialogueRunner()
        chunks: list[ChunkType] = []

        with patch(
            "swarm.runners.dialogue.execute_tools", side_effect=fake_execute_tools
        ):
            async for chunk in runner.run_streamed(
                agent=agent_a, run_context=cast(Any, run_context), max_turns=5
            ):
                chunks.append(chunk)

        # Final chunk should be the completion delimiter
        final = chunks[-1]
        assert isinstance(final, DelimiterChunk)
        assert final.delim == "complete"

        # Guardrail must have been called with Agent A, not Agent B
        assert len(guardrail_received_agent) == 1
        assert guardrail_received_agent[0].name == "AgentA"

        # Results must appear on execution_metadata
        assert final.execution_metadata is not None
        gr_results = final.execution_metadata.output_guardrail_results
        assert len(gr_results) == 1
        assert gr_results[0].tripwire_triggered is True
        assert gr_results[0].guardrail_name == "recorder"
