"""Tests for MemoryV2Client and MemoryV2."""

import pytest
from unittest.mock import AsyncMock, patch
from basion_agent.memory_v2_client import MemoryV2Client
from basion_agent.memory_v2 import MemoryV2
from basion_agent.exceptions import APIException
import aiohttp


# Helper classes for async context manager mocking
class MockResponse:
    def __init__(self, status=200, json_data=None, text_data=""):
        self.status = status
        self._json_data = json_data or {}
        self.text_data = text_data

    async def json(self):
        return self._json_data

    async def text(self):
        return self.text_data

    async def __aenter__(self):
        return self

    async def __aexit__(self, *args):
        pass


class MockSession:
    def __init__(self, response):
        self.response = response
        self.last_url = None
        self.last_json = None

    def post(self, url, json=None):
        self.last_url = url
        self.last_json = json
        return self.response


class TestMemoryV2Client:
    """Test MemoryV2Client low-level HTTP client."""

    @pytest.mark.asyncio
    async def test_initialization(self):
        """Test MemoryV2Client initialization."""
        client = MemoryV2Client(base_url="http://memory:3004")
        assert client.base_url == "http://memory:3004"
        assert client._session is None
        await client.close()

    @pytest.mark.asyncio
    async def test_base_url_normalization(self):
        """Test that trailing slashes are removed from base_url."""
        client = MemoryV2Client(base_url="http://memory:3004/")
        assert client.base_url == "http://memory:3004"
        await client.close()

    @pytest.mark.asyncio
    async def test_get_session_creates_session(self):
        """Test that _get_session creates aiohttp session."""
        client = MemoryV2Client(base_url="http://memory:3004")
        session = await client._get_session()
        assert isinstance(session, aiohttp.ClientSession)
        assert not session.closed
        await client.close()

    @pytest.mark.asyncio
    async def test_get_session_reuses_existing_session(self):
        """Test that _get_session reuses existing session."""
        client = MemoryV2Client(base_url="http://memory:3004")
        session1 = await client._get_session()
        session2 = await client._get_session()
        assert session1 is session2
        await client.close()

    @pytest.mark.asyncio
    async def test_close_session(self):
        """Test closing the client session."""
        client = MemoryV2Client(base_url="http://memory:3004")
        session = await client._get_session()
        await client.close()
        assert session.closed

    @pytest.mark.asyncio
    async def test_close_without_session(self):
        """Test closing when no session exists."""
        client = MemoryV2Client(base_url="http://memory:3004")
        await client.close()  # Should not raise

    @pytest.mark.asyncio
    async def test_ingest_success(self):
        """Test successful memory ingestion."""
        client = MemoryV2Client(base_url="http://memory:3004")

        mock_response_data = {"status": "ok", "memories": []}
        mock_response = MockResponse(status=200, json_data=mock_response_data)
        mock_session = MockSession(mock_response)

        with patch.object(client, '_get_session', return_value=mock_session):
            result = await client.ingest(
                user_id="user-123",
                content="I like cats",
                role="user",
                conversation_id="conv-456",
            )

            assert result == mock_response_data
            assert mock_session.last_url == "http://memory:3004/ingest"
            assert mock_session.last_json == {
                "user_id": "user-123",
                "role": "user",
                "content": "I like cats",
                "conversation_id": "conv-456",
            }

        await client.close()

    @pytest.mark.asyncio
    async def test_ingest_without_conversation_id(self):
        """Test ingestion without conversation_id."""
        client = MemoryV2Client(base_url="http://memory:3004")

        mock_response = MockResponse(status=200, json_data={"status": "ok"})
        mock_session = MockSession(mock_response)

        with patch.object(client, '_get_session', return_value=mock_session):
            await client.ingest(
                user_id="user-123",
                content="Hello",
                role="assistant",
            )

            assert "conversation_id" not in mock_session.last_json

        await client.close()

    @pytest.mark.asyncio
    async def test_ingest_error_response(self):
        """Test ingest handles error responses."""
        client = MemoryV2Client(base_url="http://memory:3004")

        mock_response = MockResponse(status=500, text_data="Internal Server Error")
        mock_session = MockSession(mock_response)

        with patch.object(client, '_get_session', return_value=mock_session):
            with pytest.raises(APIException, match="Memory V2 ingest failed: 500"):
                await client.ingest(
                    user_id="user-123",
                    content="test",
                )

        await client.close()

    @pytest.mark.asyncio
    async def test_ingest_network_error(self):
        """Test ingest handles network errors."""
        client = MemoryV2Client(base_url="http://memory:3004")

        class ErrorSession:
            def post(self, url, json=None):
                raise aiohttp.ClientError("Connection refused")

        with patch.object(client, '_get_session', return_value=ErrorSession()):
            with pytest.raises(APIException, match="Memory V2 ingest failed"):
                await client.ingest(
                    user_id="user-123",
                    content="test",
                )

        await client.close()

    @pytest.mark.asyncio
    async def test_search_success(self):
        """Test successful memory search."""
        client = MemoryV2Client(base_url="http://memory:3004")

        mock_response_data = [
            {"memory": "User likes cats", "score": 0.95},
            {"memory": "User has a pet cat named Luna", "score": 0.87},
        ]
        mock_response = MockResponse(status=200, json_data=mock_response_data)
        mock_session = MockSession(mock_response)

        with patch.object(client, '_get_session', return_value=mock_session):
            results = await client.search(query="pets", user_id="user-123")

            assert len(results) == 2
            assert results[0]["memory"] == "User likes cats"
            assert mock_session.last_url == "http://memory:3004/search"
            assert mock_session.last_json == {
                "query": "pets",
                "user_id": "user-123",
            }

        await client.close()

    @pytest.mark.asyncio
    async def test_search_error_response(self):
        """Test search handles error responses."""
        client = MemoryV2Client(base_url="http://memory:3004")

        mock_response = MockResponse(status=500, text_data="Internal Server Error")
        mock_session = MockSession(mock_response)

        with patch.object(client, '_get_session', return_value=mock_session):
            with pytest.raises(APIException, match="Memory V2 search failed: 500"):
                await client.search(query="test", user_id="user-123")

        await client.close()

    @pytest.mark.asyncio
    async def test_search_network_error(self):
        """Test search handles network errors."""
        client = MemoryV2Client(base_url="http://memory:3004")

        class ErrorSession:
            def post(self, url, json=None):
                raise aiohttp.ClientError("Connection refused")

        with patch.object(client, '_get_session', return_value=ErrorSession()):
            with pytest.raises(APIException, match="Memory V2 search failed"):
                await client.search(query="test", user_id="user-123")

        await client.close()


class TestMemoryV2:
    """Test MemoryV2 wrapper class."""

    @pytest.mark.asyncio
    async def test_initialization(self):
        """Test MemoryV2 initialization."""
        client = MemoryV2Client(base_url="http://memory:3004")
        memory = MemoryV2(user_id="user-1", conversation_id="conv-1", client=client)

        assert memory.user_id == "user-1"
        assert memory.conversation_id == "conv-1"
        assert memory._client is client

        await client.close()

    @pytest.mark.asyncio
    async def test_search_delegates_correctly(self):
        """Test that search delegates to client with correct user_id."""
        client = MemoryV2Client(base_url="http://memory:3004")
        memory = MemoryV2(user_id="user-123", conversation_id="conv-456", client=client)

        mock_results = [{"memory": "test result", "score": 0.9}]

        with patch.object(client, 'search', new_callable=AsyncMock) as mock_search:
            mock_search.return_value = mock_results

            results = await memory.search(query="test query")

            assert results == mock_results
            mock_search.assert_called_once_with(
                query="test query",
                user_id="user-123",
            )

        await client.close()

    @pytest.mark.asyncio
    async def test_ingest_delegates_correctly(self):
        """Test that ingest delegates to client with correct user_id and conversation_id."""
        client = MemoryV2Client(base_url="http://memory:3004")
        memory = MemoryV2(user_id="user-123", conversation_id="conv-456", client=client)

        mock_result = {"status": "ok"}

        with patch.object(client, 'ingest', new_callable=AsyncMock) as mock_ingest:
            mock_ingest.return_value = mock_result

            result = await memory.ingest(role="user", content="I like dogs")

            assert result == mock_result
            mock_ingest.assert_called_once_with(
                user_id="user-123",
                content="I like dogs",
                role="user",
                conversation_id="conv-456",
            )

        await client.close()

    @pytest.mark.asyncio
    async def test_ingest_with_different_roles(self):
        """Test ingesting with different message roles."""
        client = MemoryV2Client(base_url="http://memory:3004")
        memory = MemoryV2(user_id="user-1", conversation_id="conv-1", client=client)

        with patch.object(client, 'ingest', new_callable=AsyncMock) as mock_ingest:
            mock_ingest.return_value = {}

            for role in ["user", "assistant", "system"]:
                await memory.ingest(role=role, content=f"Message from {role}")

            assert mock_ingest.call_count == 3
            calls = mock_ingest.call_args_list
            assert calls[0][1]["role"] == "user"
            assert calls[1][1]["role"] == "assistant"
            assert calls[2][1]["role"] == "system"

        await client.close()
