"""Tests for Streamer class - Response streaming functionality."""

import pytest
from unittest.mock import MagicMock, AsyncMock, call
from basion_agent.streamer import Streamer


class TestStreamerInitialization:
    """Test Streamer initialization."""

    def test_streamer_initialization(self, mock_gateway_client, sample_message):
        """Test basic streamer initialization."""
        streamer = Streamer(
            agent_name="test-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client
        )

        assert streamer._agent_name == "test-agent"
        assert streamer._original_message == sample_message
        assert streamer._gateway_client == mock_gateway_client
        assert streamer._send_to == "user"
        assert streamer._awaiting is False
        assert streamer._finished is False

    def test_streamer_defaults(self, mock_gateway_client, sample_message):
        """Test default parameter values."""
        streamer = Streamer(
            agent_name="test-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client
        )

        assert streamer._awaiting is False
        assert streamer._send_to == "user"


class TestStreaming:
    """Test streaming content."""

    def test_stream_sends_content(self, mock_gateway_client, sample_message):
        """Test stream() sends content via gateway."""
        streamer = Streamer(
            agent_name="test-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client
        )

        streamer.stream("Hello world")

        # Verify produce was called
        mock_gateway_client.produce.assert_called_once()
        call_kwargs = mock_gateway_client.produce.call_args[1]

        assert call_kwargs["topic"] == "router.inbox"
        assert call_kwargs["body"]["content"] == "Hello world"
        assert call_kwargs["body"]["done"] is False

    def test_stream_sets_done_false(self, mock_gateway_client, sample_message):
        """Test stream() sets done=False."""
        streamer = Streamer(
            agent_name="test-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client
        )

        streamer.stream("Content")

        call_kwargs = mock_gateway_client.produce.call_args[1]
        assert call_kwargs["body"]["done"] is False

    def test_stream_uses_conversation_id(
        self, mock_gateway_client, sample_message
    ):
        """Test stream() uses conversation ID from original message."""
        streamer = Streamer(
            agent_name="test-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client
        )

        streamer.stream("Content")

        call_kwargs = mock_gateway_client.produce.call_args[1]
        assert call_kwargs["key"] == sample_message.conversation_id
        assert (
            call_kwargs["headers"]["conversationId"]
            == sample_message.conversation_id
        )

    def test_stream_sets_from_header(self, mock_gateway_client, sample_message):
        """Test stream() sets from_ header to agent name."""
        streamer = Streamer(
            agent_name="my-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client
        )

        streamer.stream("Content")

        call_kwargs = mock_gateway_client.produce.call_args[1]
        assert call_kwargs["headers"]["from_"] == "my-agent"

    def test_stream_sets_next_route(self, mock_gateway_client, sample_message):
        """Test stream() sets nextRoute header."""
        streamer = Streamer(
            agent_name="test-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client,
            send_to="agent-2"
        )

        streamer.stream("Content")

        call_kwargs = mock_gateway_client.produce.call_args[1]
        assert call_kwargs["headers"]["nextRoute"] == "agent-2"

    def test_write_alias(self, mock_gateway_client, sample_message):
        """Test write() is an alias for stream()."""
        streamer = Streamer(
            agent_name="test-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client
        )

        streamer.write("Content via write")

        # Should call produce just like stream()
        mock_gateway_client.produce.assert_called_once()
        call_kwargs = mock_gateway_client.produce.call_args[1]
        assert call_kwargs["body"]["content"] == "Content via write"


class TestFinish:
    """Test finish() method."""

    @pytest.mark.asyncio
    async def test_finish_sends_done_true(
        self, mock_gateway_client, sample_message
    ):
        """Test finish() sends done=True."""
        streamer = Streamer(
            agent_name="test-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client
        )

        await streamer.finish()

        call_kwargs = mock_gateway_client.produce.call_args[1]
        assert call_kwargs["body"]["done"] is True
        assert call_kwargs["body"]["content"] == ""

    @pytest.mark.asyncio
    async def test_finish_idempotent(self, mock_gateway_client, sample_message):
        """Test finish() can be called multiple times safely."""
        streamer = Streamer(
            agent_name="test-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client
        )

        await streamer.finish()
        await streamer.finish()  # Call again

        # Should only produce once
        assert mock_gateway_client.produce.call_count == 1

    @pytest.mark.asyncio
    async def test_finish_with_awaiting_false_no_header(
        self, mock_gateway_client, sample_message
    ):
        """Test finish() with awaiting=False does not include awaitingRoute header."""
        streamer = Streamer(
            agent_name="test-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client,
            awaiting=False
        )

        await streamer.finish()

        # The done=true produce call should NOT have awaitingRoute header
        call_kwargs = mock_gateway_client.produce.call_args[1]
        assert "awaitingRoute" not in call_kwargs["headers"]

    @pytest.mark.asyncio
    async def test_finish_with_awaiting_true_sets_awaiting_route(
        self, mock_gateway_client, sample_message, mock_conversation_client
    ):
        """Test finish() with awaiting=True sets awaiting_route via conversation client."""
        streamer = Streamer(
            agent_name="test-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client,
            conversation_client=mock_conversation_client,
            awaiting=True
        )

        await streamer.finish()

        mock_conversation_client.update_conversation.assert_called_once_with(
            sample_message.conversation_id,
            {"awaiting_route": "test-agent"}
        )


class TestContextManager:
    """Test context manager functionality."""

    @pytest.mark.asyncio
    async def test_context_manager_enter_returns_self(
        self, mock_gateway_client, sample_message
    ):
        """Test __aenter__ returns self."""
        streamer = Streamer(
            agent_name="test-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client
        )

        async with streamer as s:
            assert s is streamer

    @pytest.mark.asyncio
    async def test_context_manager_exit_calls_finish(
        self, mock_gateway_client, sample_message
    ):
        """Test __aexit__ automatically calls finish()."""
        streamer = Streamer(
            agent_name="test-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client
        )

        async with streamer as s:
            s.stream("Content")

        # finish() should have been called automatically
        # Check that done=True was sent
        calls = mock_gateway_client.produce.call_args_list
        last_call = calls[-1]
        assert last_call[1]["body"]["done"] is True

    @pytest.mark.asyncio
    async def test_context_manager_sets_awaiting_route(
        self, mock_gateway_client, sample_message, mock_conversation_client
    ):
        """Test context manager sets awaiting_route when awaiting=True."""
        streamer = Streamer(
            agent_name="test-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client,
            conversation_client=mock_conversation_client,
            awaiting=True
        )

        async with streamer as s:
            s.stream("Content")

        # Should have called update_conversation with awaiting_route
        mock_conversation_client.update_conversation.assert_called_once_with(
            sample_message.conversation_id,
            {"awaiting_route": "test-agent"}
        )


class TestFluentAPI:
    """Test fluent API methods."""

    def test_set_response_schema(self, mock_gateway_client, sample_message):
        """Test set_response_schema() returns self and stores schema."""
        streamer = Streamer(
            agent_name="test-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client
        )

        schema = {
            "type": "object",
            "properties": {"answer": {"type": "string"}}
        }
        result = streamer.set_response_schema(schema)

        # Should return self for chaining
        assert result is streamer
        # Should store schema
        assert streamer._response_schema == schema

    def test_set_message_metadata(self, mock_gateway_client, sample_message):
        """Test set_message_metadata() returns self and stores metadata."""
        streamer = Streamer(
            agent_name="test-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client
        )

        metadata = {"priority": "high", "tags": ["urgent"]}
        result = streamer.set_message_metadata(metadata)

        # Should return self for chaining
        assert result is streamer
        # Should store metadata
        assert streamer._message_metadata == metadata

    def test_fluent_chaining(self, mock_gateway_client, sample_message):
        """Test fluent API can be chained."""
        streamer = Streamer(
            agent_name="test-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client
        )

        schema = {"type": "object"}
        metadata = {"key": "value"}

        # Chain calls
        result = (
            streamer.set_response_schema(schema)
            .set_message_metadata(metadata)
        )

        assert result is streamer
        assert streamer._response_schema == schema
        assert streamer._message_metadata == metadata


class TestAwaitingAndSchema:
    """Test awaiting_route and pending_response_schema signaling via conversation client."""

    @pytest.mark.asyncio
    async def test_finish_with_schema_sets_both(
        self, mock_gateway_client, sample_message, mock_conversation_client
    ):
        """Test schema and awaiting_route are set via conversation client on finish."""
        streamer = Streamer(
            agent_name="test-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client,
            conversation_client=mock_conversation_client,
        )

        schema = {"type": "object"}
        streamer.set_response_schema(schema)
        await streamer.finish()

        # Should set both pending_response_schema and awaiting_route (auto-await for schema)
        calls = mock_conversation_client.update_conversation.call_args_list
        assert len(calls) == 2
        assert calls[0] == ((sample_message.conversation_id, {"pending_response_schema": schema}),)
        assert calls[1] == ((sample_message.conversation_id, {"awaiting_route": "test-agent"}),)

    @pytest.mark.asyncio
    async def test_finish_with_awaiting_true(
        self, mock_gateway_client, sample_message, mock_conversation_client
    ):
        """Test awaiting_route is set when awaiting=True."""
        streamer = Streamer(
            agent_name="test-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client,
            conversation_client=mock_conversation_client,
            awaiting=True,
        )

        await streamer.finish()

        mock_conversation_client.update_conversation.assert_called_once_with(
            sample_message.conversation_id,
            {"awaiting_route": "test-agent"}
        )

    @pytest.mark.asyncio
    async def test_finish_without_awaiting_no_update(
        self, mock_gateway_client, sample_message, mock_conversation_client
    ):
        """Test no conversation update when awaiting=False and no schema."""
        streamer = Streamer(
            agent_name="test-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client,
            conversation_client=mock_conversation_client,
            awaiting=False,
        )

        await streamer.finish()

        mock_conversation_client.update_conversation.assert_not_called()

    @pytest.mark.asyncio
    async def test_schema_triggers_auto_await(
        self, mock_gateway_client, sample_message, mock_conversation_client
    ):
        """Test response schema triggers auto-await unless explicitly false."""
        streamer = Streamer(
            agent_name="test-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client,
            conversation_client=mock_conversation_client,
        )

        streamer.set_response_schema({"type": "object"})
        await streamer.finish()

        # Auto-await should trigger awaiting_route
        calls = mock_conversation_client.update_conversation.call_args_list
        awaiting_call = [c for c in calls if "awaiting_route" in c[0][1]]
        assert len(awaiting_call) == 1
        assert awaiting_call[0][0][1]["awaiting_route"] == "test-agent"

    @pytest.mark.asyncio
    async def test_schema_with_explicit_no_await(
        self, mock_gateway_client, sample_message, mock_conversation_client
    ):
        """Test schema does not auto-await when awaiting is explicitly False."""
        streamer = Streamer(
            agent_name="test-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client,
            conversation_client=mock_conversation_client,
        )

        streamer.set_response_schema({"type": "object"})
        await streamer.finish(awaiting=False)

        # Should set schema but NOT awaiting_route
        calls = mock_conversation_client.update_conversation.call_args_list
        assert len(calls) == 1
        assert "pending_response_schema" in calls[0][0][1]

    @pytest.mark.asyncio
    async def test_no_update_without_conversation_client(
        self, mock_gateway_client, sample_message
    ):
        """Test no conversation update when conversation_client is not provided."""
        streamer = Streamer(
            agent_name="test-agent",
            original_message=sample_message,
            gateway_client=mock_gateway_client,
            awaiting=True,
        )

        # Should not raise even without conversation_client
        await streamer.finish()


class TestDirectConversationId:
    """Test Streamer with direct conversation_id/user_id (no original_message)."""

    def test_streamer_with_direct_ids(self, mock_gateway_client):
        """Test streamer can be created with direct conversation_id and user_id."""
        streamer = Streamer(
            agent_name="test-agent",
            gateway_client=mock_gateway_client,
            conversation_id="direct-conv-123",
            user_id="direct-user-456",
        )

        assert streamer._conversation_id == "direct-conv-123"
        assert streamer._user_id == "direct-user-456"
        assert streamer._original_message is None

    def test_streamer_direct_ids_stream_uses_correct_conversation(
        self, mock_gateway_client
    ):
        """Test stream() uses direct conversation_id in headers and key."""
        streamer = Streamer(
            agent_name="test-agent",
            gateway_client=mock_gateway_client,
            conversation_id="direct-conv-123",
            user_id="direct-user-456",
        )

        streamer.stream("Hello")

        call_kwargs = mock_gateway_client.produce.call_args[1]
        assert call_kwargs["key"] == "direct-conv-123"
        assert call_kwargs["headers"]["conversationId"] == "direct-conv-123"
        assert call_kwargs["headers"]["userId"] == "direct-user-456"

    def test_streamer_direct_ids_uses_agent_name_in_headers(
        self, mock_gateway_client
    ):
        """Test from_ header uses agent name with direct IDs."""
        streamer = Streamer(
            agent_name="my-agent",
            gateway_client=mock_gateway_client,
            conversation_id="conv-1",
            user_id="user-1",
        )

        streamer.stream("Content")

        call_kwargs = mock_gateway_client.produce.call_args[1]
        assert call_kwargs["headers"]["from_"] == "my-agent"

    def test_streamer_requires_either_message_or_direct_ids(
        self, mock_gateway_client
    ):
        """Test streamer raises ValueError if neither message nor direct IDs provided."""
        with pytest.raises(ValueError, match="Either original_message or both"):
            Streamer(
                agent_name="test-agent",
                gateway_client=mock_gateway_client,
            )

    def test_streamer_requires_both_conversation_id_and_user_id(
        self, mock_gateway_client
    ):
        """Test streamer raises ValueError if only conversation_id provided."""
        with pytest.raises(ValueError, match="Either original_message or both"):
            Streamer(
                agent_name="test-agent",
                gateway_client=mock_gateway_client,
                conversation_id="conv-1",
            )

    def test_streamer_requires_both_user_id_and_conversation_id(
        self, mock_gateway_client
    ):
        """Test streamer raises ValueError if only user_id provided."""
        with pytest.raises(ValueError, match="Either original_message or both"):
            Streamer(
                agent_name="test-agent",
                gateway_client=mock_gateway_client,
                user_id="user-1",
            )

    def test_streamer_original_message_takes_precedence(
        self, mock_gateway_client, sample_message
    ):
        """Test original_message IDs take precedence over direct IDs."""
        streamer = Streamer(
            agent_name="test-agent",
            gateway_client=mock_gateway_client,
            original_message=sample_message,
            conversation_id="should-be-ignored",
            user_id="should-be-ignored",
        )

        # Should use IDs from original_message, not direct ones
        assert streamer._conversation_id == sample_message.conversation_id
        assert streamer._user_id == sample_message.user_id

    @pytest.mark.asyncio
    async def test_finish_with_direct_ids_sets_awaiting_route(
        self, mock_gateway_client, mock_conversation_client
    ):
        """Test finish() sets awaiting_route via conversation client with direct IDs."""
        streamer = Streamer(
            agent_name="test-agent",
            gateway_client=mock_gateway_client,
            conversation_client=mock_conversation_client,
            conversation_id="direct-conv-123",
            user_id="direct-user-456",
            awaiting=True,
        )

        await streamer.finish()

        mock_conversation_client.update_conversation.assert_called_once_with(
            "direct-conv-123",
            {"awaiting_route": "test-agent"}
        )

    @pytest.mark.asyncio
    async def test_context_manager_with_direct_ids(
        self, mock_gateway_client
    ):
        """Test context manager works with direct IDs."""
        streamer = Streamer(
            agent_name="test-agent",
            gateway_client=mock_gateway_client,
            conversation_id="direct-conv-123",
            user_id="direct-user-456",
        )

        async with streamer as s:
            s.stream("Content via context manager")

        # Verify done was sent
        calls = mock_gateway_client.produce.call_args_list
        last_call = calls[-1]
        assert last_call[1]["body"]["done"] is True
