"""Tests for BaseAgent streaming functionality."""

from typing import Callable, Generator
from unittest.mock import Mock

import pytest

from thetable_poc.config.llm_provider import LLMProvider
from thetable_poc.core.agents.base_agent import BaseAgent
from thetable_poc.core.participant import ParticipantConfig
from thetable_poc.core.roles import Role, ParticipantType


class MockStreamingProvider(LLMProvider):
    """Mock LLM provider for testing streaming."""

    def __init__(self, responses: list[str] = None):
        self.responses = responses or ["Hello", " world", "!"]
        self.call_count = 0

    def chat(self, system: str, user: str, temperature: float = 0.7) -> str:
        """Non-streaming chat."""
        self.call_count += 1
        return "".join(self.responses)

    def chat_stream(
        self, system: str, user: str, temperature: float = 0.7
    ) -> Generator[str, None, None]:
        """Streaming chat."""
        self.call_count += 1
        for chunk in self.responses:
            yield chunk


class TestBaseAgentStreaming:
    """Test BaseAgent streaming integration."""

    def test_speak_stream_calls_typing_callback(self):
        """Test that speak_stream calls typing callback for each chunk."""
        # Setup
        config = ParticipantConfig(
            name="TestAgent",
            role=Role.HOST,
            participant_type=ParticipantType.AI,
        )
        provider = MockStreamingProvider(["Hello", " ", "world"])
        llm_settings = {"temperature": 0.7}
        agent = BaseAgent(config, provider=provider, llm_settings=llm_settings)

        # Track callback invocations
        callback_chunks = []

        def typing_callback(chunk: str):
            callback_chunks.append(chunk)

        # Execute
        result = agent.speak_stream(context={}, typing_callback=typing_callback)

        # Verify
        assert callback_chunks == ["Hello", " ", "world"]
        assert result == "Hello world"

    def test_speak_stream_returns_full_response(self):
        """Test that speak_stream returns the full accumulated response."""
        config = ParticipantConfig(
            name="TestAgent",
            role=Role.HOST,
            participant_type=ParticipantType.AI,
        )
        provider = MockStreamingProvider(["Test", " ", "response"])
        llm_settings = {"temperature": 0.7}
        agent = BaseAgent(config, provider=provider, llm_settings=llm_settings)

        result = agent.speak_stream(context={})

        assert result == "Test response"

    def test_speak_stream_accumulates_chunks(self):
        """Test that chunks are properly accumulated."""
        config = ParticipantConfig(
            name="TestAgent",
            role=Role.HOST,
            participant_type=ParticipantType.AI,
        )
        provider = MockStreamingProvider(["A", "B", "C", "D"])
        llm_settings = {"temperature": 0.7}
        agent = BaseAgent(config, provider=provider, llm_settings=llm_settings)

        result = agent.speak_stream(context={})

        assert result == "ABCD"

    def test_speak_stream_works_without_callback(self):
        """Test that speak_stream works even without a callback."""
        config = ParticipantConfig(
            name="TestAgent",
            role=Role.HOST,
            participant_type=ParticipantType.AI,
        )
        provider = MockStreamingProvider(["Hello", " world"])
        llm_settings = {"temperature": 0.7}
        agent = BaseAgent(config, provider=provider, llm_settings=llm_settings)

        # Should not raise error without callback
        result = agent.speak_stream(context={})

        assert result == "Hello world"

    def test_speak_stream_handles_empty_chunks(self):
        """Test that empty chunks are handled gracefully."""
        config = ParticipantConfig(
            name="TestAgent",
            role=Role.HOST,
            participant_type=ParticipantType.AI,
        )
        provider = MockStreamingProvider(["Hello", "", " ", "", "world"])
        llm_settings = {"temperature": 0.7}
        agent = BaseAgent(config, provider=provider, llm_settings=llm_settings)

        callback_chunks = []

        def typing_callback(chunk: str):
            callback_chunks.append(chunk)

        result = agent.speak_stream(context={}, typing_callback=typing_callback)

        # Empty chunks should still be passed to callback
        assert callback_chunks == ["Hello", "", " ", "", "world"]
        assert result == "Hello world"

    def test_speak_stream_uses_same_prompt_as_speak(self):
        """Test that speak_stream uses the same prompt building logic."""
        config = ParticipantConfig(
            name="TestAgent",
            role=Role.HOST,
            participant_type=ParticipantType.AI,
        )
        provider = MockStreamingProvider()
        llm_settings = {"temperature": 0.7}
        agent = BaseAgent(config, provider=provider, llm_settings=llm_settings)

        context = {
            "phase": "discussion",
            "agenda": "Test agenda",
            "recent_messages": [
                {"sender": "User1", "message": "Message 1"},
                {"sender": "User2", "message": "Message 2"},
            ],
            "question": "What do you think?",
        }

        # Both methods should work with the same context
        result1 = agent.speak(context)
        result2 = agent.speak_stream(context)

        # Should return same content (though may be constructed differently)
        assert len(result1) > 0
        assert len(result2) > 0
