"""Simplified tests for DSPy streaming WebSocket event emission.

Tests verify the core behavior: when tokens arrive from DSPy, StreamingOutputEvent
is created with correct parameters and broadcast via WebSocket.

Specification: docs/specs/005-streaming-output-fix
"""

from unittest.mock import AsyncMock, Mock
from uuid import uuid4

import pytest

from flock.components.server.models.events import StreamingOutputEvent
from flock.engines.dspy_engine import DSPyEngine


class TestWebSocketManagerRetrieval:
    """Test that engine retrieves WebSocket manager from context."""

    def test_gets_manager_from_orchestrator_collector(self):
        """Test WebSocket manager is retrieved via orchestrator->collector chain."""
        # Arrange
        DSPyEngine()

        # Create mock chain: ctx -> orchestrator -> collector -> ws_manager
        mock_ws_manager = Mock()
        mock_collector = Mock()
        mock_collector._websocket_manager = mock_ws_manager
        mock_orchestrator = Mock()
        mock_orchestrator._dashboard_collector = mock_collector

        ctx = Mock()
        ctx.orchestrator = mock_orchestrator

        # Act - Simulate what happens in _execute_streaming
        ws_manager = None
        if ctx:
            orchestrator = getattr(ctx, "orchestrator", None)
            if orchestrator:
                collector = getattr(orchestrator, "_dashboard_collector", None)
                if collector:
                    ws_manager = getattr(collector, "_websocket_manager", None)

        # Assert
        assert ws_manager is mock_ws_manager

    def test_graceful_degradation_without_orchestrator(self):
        """Test that missing orchestrator doesn't crash."""
        # Arrange
        ctx = Mock(spec=[])  # No orchestrator attribute

        # Act
        ws_manager = None
        if ctx:
            orchestrator = getattr(ctx, "orchestrator", None)
            if orchestrator:
                collector = getattr(orchestrator, "_dashboard_collector", None)
                if collector:
                    ws_manager = getattr(collector, "_websocket_manager", None)

        # Assert - Should be None, not crash
        assert ws_manager is None

    def test_graceful_degradation_without_collector(self):
        """Test that missing collector doesn't crash."""
        # Arrange
        ctx = Mock()
        ctx.orchestrator = Mock(spec=[])  # No _dashboard_collector

        # Act
        ws_manager = None
        if ctx:
            orchestrator = getattr(ctx, "orchestrator", None)
            if orchestrator:
                collector = getattr(orchestrator, "_dashboard_collector", None)
                if collector:
                    ws_manager = getattr(collector, "_websocket_manager", None)

        # Assert
        assert ws_manager is None


class TestStreamingOutputEventCreation:
    """Test that StreamingOutputEvent is created with correct parameters."""

    def test_status_message_creates_log_event(self):
        """Test that status messages create events with output_type='log'."""
        # Arrange
        ctx = Mock()
        ctx.correlation_id = str(uuid4())
        ctx.task_id = "task-123"
        agent = Mock()
        agent.name = "test_agent"

        # Act - Simulate event creation for StatusMessage
        event = StreamingOutputEvent(
            correlation_id=str(ctx.correlation_id),
            agent_name=agent.name,
            run_id=ctx.task_id,
            output_type="log",
            content="Processing...",
            sequence=0,
            is_final=False,
        )

        # Assert
        assert event.output_type == "log"
        assert event.content == "Processing..."
        assert event.agent_name == "test_agent"
        assert event.run_id == "task-123"
        assert event.sequence == 0
        assert event.is_final is False

    def test_stream_response_creates_llm_token_event(self):
        """Test that stream responses create events with output_type='llm_token'."""
        # Arrange
        ctx = Mock()
        ctx.correlation_id = str(uuid4())
        ctx.task_id = "task-123"
        agent = Mock()
        agent.name = "test_agent"

        # Act - Simulate event creation for StreamResponse
        event = StreamingOutputEvent(
            correlation_id=str(ctx.correlation_id),
            agent_name=agent.name,
            run_id=ctx.task_id,
            output_type="llm_token",
            content="Hello ",
            sequence=5,
            is_final=False,
        )

        # Assert
        assert event.output_type == "llm_token"
        assert event.content == "Hello "
        assert event.sequence == 5

    def test_final_event_has_is_final_true(self):
        """Test that final events are marked with is_final=True."""
        # Arrange
        ctx = Mock()
        ctx.correlation_id = str(uuid4())
        ctx.task_id = "task-123"
        agent = Mock()
        agent.name = "test_agent"

        # Act - Simulate final event creation
        event = StreamingOutputEvent(
            correlation_id=str(ctx.correlation_id),
            agent_name=agent.name,
            run_id=ctx.task_id,
            output_type="log",
            content="--- End of output ---",
            sequence=42,
            is_final=True,
        )

        # Assert
        assert event.is_final is True
        assert event.content == "--- End of output ---"

    def test_event_includes_correlation_id(self):
        """Test that events include correlation_id from context."""
        # Arrange
        correlation_id = str(uuid4())
        ctx = Mock()
        ctx.correlation_id = correlation_id
        ctx.task_id = "task-123"
        agent = Mock()
        agent.name = "test_agent"

        # Act
        event = StreamingOutputEvent(
            correlation_id=str(ctx.correlation_id),
            agent_name=agent.name,
            run_id=ctx.task_id,
            output_type="llm_token",
            content="token",
            sequence=0,
            is_final=False,
        )

        # Assert
        assert event.correlation_id == str(correlation_id)

    def test_timestamp_is_auto_generated(self):
        """Test that timestamp is automatically generated."""
        # Arrange
        ctx = Mock()
        ctx.correlation_id = str(uuid4())
        ctx.task_id = "task-123"
        agent = Mock()
        agent.name = "test_agent"

        # Act
        event = StreamingOutputEvent(
            correlation_id=str(ctx.correlation_id),
            agent_name=agent.name,
            run_id=ctx.task_id,
            output_type="llm_token",
            content="token",
            sequence=0,
            is_final=False,
        )

        # Assert - timestamp should exist and be ISO format
        assert event.timestamp is not None
        assert isinstance(event.timestamp, str)
        assert "T" in event.timestamp  # ISO format contains T


class TestSequenceOrdering:
    """Test that sequence numbers work correctly."""

    def test_sequence_increments_monotonically(self):
        """Test that sequence numbers can increment."""
        # Arrange
        ctx = Mock()
        ctx.correlation_id = str(uuid4())
        ctx.task_id = "task-123"
        agent = Mock()
        agent.name = "test_agent"

        sequence = 0

        # Act - Create multiple events
        events = []
        for i in range(5):
            event = StreamingOutputEvent(
                correlation_id=str(ctx.correlation_id),
                agent_name=agent.name,
                run_id=ctx.task_id,
                output_type="llm_token",
                content=f"token{i}",
                sequence=sequence,
                is_final=False,
            )
            events.append(event)
            sequence += 1

        # Assert - Sequences should be 0, 1, 2, 3, 4
        assert [e.sequence for e in events] == [0, 1, 2, 3, 4]


class TestEventSerialization:
    """Test that events serialize correctly for WebSocket."""

    def test_event_serializes_to_json(self):
        """Test that StreamingOutputEvent serializes to valid JSON."""
        # Arrange
        event = StreamingOutputEvent(
            correlation_id=str(uuid4()),
            agent_name="test_agent",
            run_id="task-123",
            output_type="llm_token",
            content="Hello",
            sequence=0,
            is_final=False,
        )

        # Act
        json_str = event.model_dump_json()

        # Assert
        assert isinstance(json_str, str)
        assert "test_agent" in json_str
        assert "llm_token" in json_str
        assert "Hello" in json_str

    def test_event_dict_has_required_fields(self):
        """Test that event dict contains all required fields."""
        # Arrange
        event = StreamingOutputEvent(
            correlation_id=str(uuid4()),
            agent_name="test_agent",
            run_id="task-123",
            output_type="llm_token",
            content="Hello",
            sequence=0,
            is_final=False,
        )

        # Act
        event_dict = event.model_dump()

        # Assert - Check all required fields exist
        assert "correlation_id" in event_dict
        assert "timestamp" in event_dict
        assert "agent_name" in event_dict
        assert "run_id" in event_dict
        assert "output_type" in event_dict
        assert "content" in event_dict
        assert "sequence" in event_dict
        assert "is_final" in event_dict


class TestStreamingDisabledInTests:
    """Test that streaming is auto-disabled in pytest."""

    def test_dspy_engine_detects_pytest(self):
        """Test that engine detects pytest environment."""
        import sys

        # Assert - pytest should be in sys.modules during test
        assert "pytest" in sys.modules

    def test_engine_defaults_to_no_streaming_in_tests(self):
        """Test that DSPyEngine defaults stream=False in pytest."""
        # Act
        engine = DSPyEngine()

        # Assert - Should auto-detect pytest and disable streaming
        assert engine.stream is False

    def test_engine_respects_explicit_stream_true(self):
        """Test that explicit stream=True overrides auto-detection."""
        # Act
        engine = DSPyEngine(stream=True)

        # Assert
        assert engine.stream is True


@pytest.mark.asyncio
async def test_websocket_broadcast_called_with_event():
    """Integration test: verify broadcast is called with StreamingOutputEvent."""
    # Arrange
    mock_ws_manager = AsyncMock()

    # Create a StreamingOutputEvent
    event = StreamingOutputEvent(
        correlation_id=str(uuid4()),
        agent_name="test_agent",
        run_id="task-123",
        output_type="llm_token",
        content="token",
        sequence=0,
        is_final=False,
    )

    # Act - Simulate what happens in dspy_engine
    try:
        await mock_ws_manager.broadcast(event)
    except Exception as e:
        # Should not raise
        pytest.fail(f"Broadcast raised exception: {e}")

    # Assert
    assert mock_ws_manager.broadcast.called
    assert mock_ws_manager.broadcast.call_count == 1

    # Verify the event passed to broadcast
    call_args = mock_ws_manager.broadcast.call_args
    assert call_args[0][0] == event


@pytest.mark.asyncio
async def test_broadcast_error_is_caught_gracefully():
    """Test that broadcast errors don't crash the engine."""
    # Arrange
    mock_ws_manager = AsyncMock()
    mock_ws_manager.broadcast.side_effect = Exception("WebSocket disconnected")

    event = StreamingOutputEvent(
        correlation_id=str(uuid4()),
        agent_name="test_agent",
        run_id="task-123",
        output_type="llm_token",
        content="token",
        sequence=0,
        is_final=False,
    )

    # Act - Simulate error handling in dspy_engine
    error_caught = False
    try:
        await mock_ws_manager.broadcast(event)
    except Exception:
        error_caught = True

    # Assert - Error should be caught (in real code, wrapped in try/except)
    assert error_caught  # This simulates that exception would be raised
    # In actual code, this is caught with try/except and logged as warning
