"""Tests for the Anthropic client wrapper."""

from unittest.mock import AsyncMock, MagicMock

import pytest

from aal_sdk.anthropic import _StreamWrapper, wrap_anthropic_client
from aal_sdk.session import AALSession


def _make_session() -> AALSession:
    """Create a test session with a mock push function."""
    return AALSession("sess-1", "agent-1", MagicMock())


def _make_tool_use_block(name: str = "get_weather", input_data: dict | None = None):
    """Create a mock tool_use content block."""
    block = MagicMock()
    block.type = "tool_use"
    block.name = name
    block.input = input_data or {"city": "London"}
    return block


def _make_text_block(text: str = "Hello"):
    """Create a mock text content block."""
    block = MagicMock()
    block.type = "text"
    block.text = text
    return block


def _make_client(response=None):
    """Create a mock Anthropic client."""
    client = MagicMock()
    client.messages = MagicMock()
    client.messages.create = AsyncMock(return_value=response)
    return client


async def test_non_streaming_logs_tool_use_blocks():
    """Test that non-streaming response logs tool_use blocks."""
    response = MagicMock()
    response.content = [
        _make_text_block("I'll check the weather"),
        _make_tool_use_block("get_weather", {"city": "NYC"}),
    ]

    client = _make_client(response)
    session = _make_session()

    wrapped = wrap_anthropic_client(client, session)
    result = await wrapped.messages.create(model="claude-3", messages=[])

    assert result is response
    session._push_fn.assert_called_once()
    event = session._push_fn.call_args[0][0]
    assert event["action_type"] == "TOOL_CALL"
    assert event["tool_name"] == "get_weather"
    assert event["parameters"] == {"city": "NYC"}
    assert event["status"] == "SUCCESS"


async def test_non_streaming_no_tools_is_noop():
    """Test that non-streaming response with no tool_use is a no-op."""
    response = MagicMock()
    response.content = [_make_text_block("Just text")]

    client = _make_client(response)
    session = _make_session()

    wrapped = wrap_anthropic_client(client, session)
    result = await wrapped.messages.create(model="claude-3", messages=[])

    assert result is response
    session._push_fn.assert_not_called()


async def test_streaming_accumulates_tool_use_blocks():
    """Test that streaming response accumulates and logs tool_use blocks."""
    session = _make_session()

    # Create mock stream events
    block_start = MagicMock()
    block_start.type = "content_block_start"
    block_start.index = 0
    cb = MagicMock()
    cb.type = "tool_use"
    cb.name = "search"
    block_start.content_block = cb

    delta1 = MagicMock()
    delta1.type = "content_block_delta"
    delta1.index = 0
    d1 = MagicMock()
    d1.type = "input_json_delta"
    d1.partial_json = '{"query": '
    delta1.delta = d1

    delta2 = MagicMock()
    delta2.type = "content_block_delta"
    delta2.index = 0
    d2 = MagicMock()
    d2.type = "input_json_delta"
    d2.partial_json = '"hello"}'
    delta2.delta = d2

    end_event = MagicMock()
    end_event.type = "message_stop"

    # Create an async iterable from the events
    async def mock_stream():
        for event in [block_start, delta1, delta2, end_event]:
            yield event

    stream_wrapper = _StreamWrapper(mock_stream(), session, 0.0)

    events_received = []
    async for event in stream_wrapper:
        events_received.append(event)

    assert len(events_received) == 4
    session._push_fn.assert_called_once()
    event = session._push_fn.call_args[0][0]
    assert event["tool_name"] == "search"
    assert event["parameters"] == {"query": "hello"}
    assert event["status"] == "SUCCESS"


async def test_stream_error_sets_status_to_error():
    """Test that stream error sets logged status to ERROR."""
    session = _make_session()

    block_start = MagicMock()
    block_start.type = "content_block_start"
    block_start.index = 0
    cb = MagicMock()
    cb.type = "tool_use"
    cb.name = "failing_tool"
    block_start.content_block = cb

    async def error_stream():
        yield block_start
        raise ConnectionError("Stream interrupted")

    stream_wrapper = _StreamWrapper(error_stream(), session, 0.0)

    with pytest.raises(ConnectionError, match="Stream interrupted"):
        async for _ in stream_wrapper:
            pass

    session._push_fn.assert_called_once()
    event = session._push_fn.call_args[0][0]
    assert event["status"] == "ERROR"
    assert event["tool_name"] == "failing_tool"


async def test_original_response_returned_unchanged():
    """Test that the original response object is returned unchanged."""
    response = MagicMock()
    response.content = [
        _make_tool_use_block("tool1"),
        _make_tool_use_block("tool2"),
    ]

    client = _make_client(response)
    session = _make_session()

    wrapped = wrap_anthropic_client(client, session)
    result = await wrapped.messages.create(model="claude-3", messages=[])

    # The exact same response object should be returned
    assert result is response
    # Both tool_use blocks should be logged
    assert session._push_fn.call_count == 2
