"""Tests for EvoScientist/stream/events.py helpers."""

import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock

from langchain_core.messages import AIMessageChunk

from EvoScientist.stream.events import _extract_tool_content, stream_agent_events


class TestExtractToolContent:
    """Verify _extract_tool_content handles image and text ToolMessages."""

    def test_image_via_additional_kwargs(self):
        """Image ToolMessages with read_file_media_type return summary."""
        msg = SimpleNamespace(
            content=[{"type": "image", "base64": "abc123..."}],
            additional_kwargs={
                "read_file_media_type": "image/png",
                "read_file_path": "/chart.png",
            },
            name="read_file",
        )
        content, is_image = _extract_tool_content(msg)
        assert is_image is True
        assert "chart.png" in content
        assert "image/png" in content
        # Must NOT contain base64 data
        assert "abc123" not in content

    def test_image_via_list_content_blocks(self):
        """Image content blocks without metadata are still detected."""
        msg = SimpleNamespace(
            content=[
                {"type": "text", "text": "Image: chart.png"},
                {"type": "image", "base64": "iVBORw0KGgo..."},
            ],
            additional_kwargs={},
            name="read_file",
        )
        content, is_image = _extract_tool_content(msg)
        assert is_image is True
        assert "iVBORw0KGgo" not in content

    def test_normal_text_passthrough(self):
        """Normal text content passes through unchanged."""
        msg = SimpleNamespace(
            content="File written successfully to /output.txt",
            additional_kwargs={},
            name="write_file",
        )
        content, is_image = _extract_tool_content(msg)
        assert is_image is False
        assert content == "File written successfully to /output.txt"

    def test_empty_content(self):
        """Empty content returns empty string."""
        msg = SimpleNamespace(
            content="",
            additional_kwargs={},
            name="read_file",
        )
        content, is_image = _extract_tool_content(msg)
        assert is_image is False
        assert content == ""

    def test_list_text_blocks(self):
        """List of text blocks are joined."""
        msg = SimpleNamespace(
            content=[
                {"type": "text", "text": "Line 1"},
                {"type": "text", "text": "Line 2"},
            ],
            additional_kwargs={},
            name="read_file",
        )
        content, is_image = _extract_tool_content(msg)
        assert is_image is False
        assert "Line 1" in content
        assert "Line 2" in content

    def test_no_additional_kwargs_attr(self):
        """Messages without additional_kwargs attribute are handled."""
        msg = SimpleNamespace(
            content="some result",
            name="execute",
        )
        content, is_image = _extract_tool_content(msg)
        assert is_image is False
        assert content == "some result"


# =============================================================================
# Multi-mode streaming chunk unpacking
# =============================================================================


def _make_ai_chunk(content: str = "hello", **kwargs):
    """Create a minimal AIMessageChunk for testing."""
    return AIMessageChunk(content=content, **kwargs)


def _collect_events(agent, message="hi", thread_id="t1"):
    """Collect all events from stream_agent_events synchronously."""
    async def _run():
        events = []
        async for ev in stream_agent_events(agent, message, thread_id):
            events.append(ev)
        return events
    loop = asyncio.new_event_loop()
    try:
        return loop.run_until_complete(_run())
    finally:
        loop.close()


async def _async_iter(items):
    """Create an async iterator from a list."""
    for item in items:
        yield item


class TestMultiModeChunkUnpacking:
    """Test 3-tuple (multi-mode) and 2-tuple (single-mode) chunk handling."""

    def test_3tuple_chunk_unpacking(self):
        """Multi-mode yields 3-tuples (namespace, mode, data); messages are processed."""
        chunk = _make_ai_chunk("hello world")
        mock_agent = AsyncMock()
        mock_agent.astream = MagicMock(return_value=_async_iter([
            ((), "messages", (chunk, {})),
        ]))
        events = _collect_events(mock_agent)
        text_events = [e for e in events if e.get("type") == "text"]
        assert len(text_events) == 1
        assert text_events[0]["content"] == "hello world"

    def test_2tuple_fallback(self):
        """Single-mode yields 2-tuples; should still work."""
        chunk = _make_ai_chunk("fallback")
        mock_agent = AsyncMock()
        mock_agent.astream = MagicMock(return_value=_async_iter([
            ((), (chunk, {})),
        ]))
        events = _collect_events(mock_agent)
        text_events = [e for e in events if e.get("type") == "text"]
        assert len(text_events) == 1
        assert text_events[0]["content"] == "fallback"

    def test_updates_mode_graceful_skip(self):
        """Updates mode chunks are skipped without error."""
        chunk = _make_ai_chunk("should appear")
        mock_agent = AsyncMock()
        mock_agent.astream = MagicMock(return_value=_async_iter([
            ((), "updates", {"some": "state"}),
            ((), "messages", (chunk, {})),
        ]))
        events = _collect_events(mock_agent)
        text_events = [e for e in events if e.get("type") == "text"]
        assert len(text_events) == 1
        assert text_events[0]["content"] == "should appear"

    def test_summarization_filtered(self):
        """Chunks with lc_source=summarization metadata are filtered out."""
        chunk_real = _make_ai_chunk("real content")
        chunk_synth = _make_ai_chunk("synthetic summary")
        mock_agent = AsyncMock()
        mock_agent.astream = MagicMock(return_value=_async_iter([
            ((), "messages", (chunk_synth, {"lc_source": "summarization"})),
            ((), "messages", (chunk_real, {})),
        ]))
        events = _collect_events(mock_agent)
        text_events = [e for e in events if e.get("type") == "text"]
        assert len(text_events) == 1
        assert text_events[0]["content"] == "real content"


class TestUsageStatsExtraction:
    """Test token usage extraction from AIMessageChunk."""

    def test_usage_metadata_emitted(self):
        """AIMessageChunk with usage_metadata emits usage_stats event."""
        chunk = _make_ai_chunk("hi", usage_metadata={"input_tokens": 100, "output_tokens": 50, "total_tokens": 150})
        mock_agent = AsyncMock()
        mock_agent.astream = MagicMock(return_value=_async_iter([
            ((), "messages", (chunk, {})),
        ]))
        events = _collect_events(mock_agent)
        usage_events = [e for e in events if e.get("type") == "usage_stats"]
        assert len(usage_events) == 1
        assert usage_events[0]["input_tokens"] == 100
        assert usage_events[0]["output_tokens"] == 50

    def test_no_usage_metadata_no_event(self):
        """AIMessageChunk without usage_metadata does not emit usage_stats."""
        chunk = _make_ai_chunk("hi")
        mock_agent = AsyncMock()
        mock_agent.astream = MagicMock(return_value=_async_iter([
            ((), "messages", (chunk, {})),
        ]))
        events = _collect_events(mock_agent)
        usage_events = [e for e in events if e.get("type") == "usage_stats"]
        assert len(usage_events) == 0

    def test_zero_tokens_not_emitted(self):
        """Zero input and output tokens should not emit usage_stats."""
        chunk = _make_ai_chunk("hi", usage_metadata={"input_tokens": 0, "output_tokens": 0, "total_tokens": 0})
        mock_agent = AsyncMock()
        mock_agent.astream = MagicMock(return_value=_async_iter([
            ((), "messages", (chunk, {})),
        ]))
        events = _collect_events(mock_agent)
        usage_events = [e for e in events if e.get("type") == "usage_stats"]
        assert len(usage_events) == 0
