"""
Tests for MCALMemory - AutoGen Memory interface implementation.

Tests cover:
- Initialization and configuration
- add() method
- query() method  
- update_context() method
- clear() method
- close() method
- TTL and expiration
- Thread safety
- Error handling
"""

import asyncio
import pytest
import time
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
from unittest.mock import MagicMock, AsyncMock, patch


# ============================================================================
# Mock AutoGen types for testing without autogen-core installed
# ============================================================================

class MockMemoryMimeType:
    """Mock of autogen_core.memory.MemoryMimeType."""
    TEXT = "text/plain"
    JSON = "application/json"
    MARKDOWN = "text/markdown"
    BINARY = "application/octet-stream"


class MockMemoryContent:
    """Mock of autogen_core.memory.MemoryContent."""
    def __init__(
        self,
        content: Any,
        mime_type: str = "text/plain",
        metadata: Optional[Dict[str, Any]] = None,
    ):
        self.content = content
        self.mime_type = mime_type
        self.metadata = metadata or {}


class MockMemoryQueryResult:
    """Mock of autogen_core.memory.MemoryQueryResult."""
    def __init__(self, results: List[MockMemoryContent]):
        self.results = results


class MockUpdateContextResult:
    """Mock of autogen_core.memory.UpdateContextResult."""
    def __init__(self, memories: MockMemoryQueryResult):
        self.memories = memories


class MockChatCompletionContext:
    """Mock of autogen_core.model_context.ChatCompletionContext."""
    def __init__(self, messages: Optional[List[Any]] = None):
        self._messages = messages or []
    
    async def get_messages(self) -> List[Any]:
        return self._messages
    
    async def add_message(self, message: Any) -> None:
        self._messages.append(message)


class MockSystemMessage:
    """Mock of autogen_core.models.SystemMessage."""
    def __init__(self, content: str):
        self.content = content


class MockUserMessage:
    """Mock of a user message."""
    def __init__(self, content: str):
        self.content = content


class MockCancellationToken:
    """Mock of autogen_core.CancellationToken."""
    def __init__(self, cancelled: bool = False):
        self.cancelled = cancelled


# ============================================================================
# Mock MCAL for testing
# ============================================================================

class MockMCAL:
    """Mock MCAL instance for testing."""
    
    def __init__(self, goal: str = "Test goal"):
        self.goal = goal
        self._memories: List[Dict[str, Any]] = []
    
    async def add(
        self,
        messages: List[Dict[str, str]],
        user_id: str = "default",
        metadata: Optional[Dict[str, Any]] = None,
    ) -> MagicMock:
        """Mock async add method matching MCAL.add() signature."""
        for msg in messages:
            self._memories.append({
                "content": msg.get("content", ""),
                "metadata": metadata or {},
                "user_id": user_id,
            })
        return MagicMock()
    
    async def search(
        self,
        query: str,
        user_id: str = "default",
        limit: int = 10,
        include_goals: bool = True,
        include_decisions: bool = True,
        retrieval_config: Any = None,
    ) -> MagicMock:
        """Mock async search returning SearchResult-like object."""
        results = []
        query_lower = query.lower()
        
        for mem in self._memories:
            if user_id != "default" and mem.get("user_id") != user_id:
                continue
            
            content = mem.get("content", "")
            if query_lower in content.lower():
                mock_entry = MagicMock()
                mock_entry.id = f"mem_{len(results)}"
                mock_entry.content = content
                mock_entry.score = 0.8
                mock_entry.metadata = mem.get("metadata", {})
                mock_entry.created_at = None
                mock_entry.updated_at = None
                results.append(mock_entry)
        
        search_result = MagicMock()
        search_result.memories = results[:limit]
        search_result.active_goals = []
        search_result.relevant_decisions = []
        search_result.context = None
        return search_result
    
    def clear_user(self, user_id: str) -> None:
        """Mock clear_user method."""
        self._memories = [
            m for m in self._memories
            if m.get("user_id") != user_id
        ]
    
    def reset(self) -> None:
        """Mock reset method."""
        self._memories = []


# ============================================================================
# Patch AutoGen imports for testing
# ============================================================================

@pytest.fixture(autouse=True)
def mock_autogen_imports(monkeypatch):
    """Patch AutoGen imports with mocks."""
    import sys
    
    # Patch the _compat module
    monkeypatch.setattr(
        "mcal_autogen._compat.AUTOGEN_CORE_AVAILABLE",
        True
    )
    monkeypatch.setattr(
        "mcal_autogen._compat.check_autogen_core",
        lambda: None
    )
    monkeypatch.setattr(
        "mcal_autogen._compat.MemoryContent",
        MockMemoryContent
    )
    monkeypatch.setattr(
        "mcal_autogen._compat.MemoryQueryResult",
        MockMemoryQueryResult
    )
    monkeypatch.setattr(
        "mcal_autogen._compat.UpdateContextResult",
        MockUpdateContextResult
    )
    monkeypatch.setattr(
        "mcal_autogen._compat.MemoryMimeType",
        MockMemoryMimeType
    )
    monkeypatch.setattr(
        "mcal_autogen._compat.SystemMessage",
        MockSystemMessage
    )
    
    # Also patch in the memory module where they're imported
    monkeypatch.setattr(
        "mcal_autogen.memory.MemoryContent",
        MockMemoryContent
    )
    monkeypatch.setattr(
        "mcal_autogen.memory.MemoryQueryResult",
        MockMemoryQueryResult
    )
    monkeypatch.setattr(
        "mcal_autogen.memory.UpdateContextResult",
        MockUpdateContextResult
    )
    monkeypatch.setattr(
        "mcal_autogen.memory.MemoryMimeType",
        MockMemoryMimeType
    )
    monkeypatch.setattr(
        "mcal_autogen.memory.SystemMessage",
        MockSystemMessage
    )


# ============================================================================
# Test Fixtures
# ============================================================================

@pytest.fixture
def mock_mcal():
    """Create a mock MCAL instance."""
    return MockMCAL(goal="Build test system")


@pytest.fixture
def memory(mock_mcal):
    """Create a MCALMemory instance for testing."""
    from mcal_autogen.memory import MCALMemory
    return MCALMemory(
        mcal=mock_mcal,
        user_id="test_user",
        name="test_memory",
    )


@pytest.fixture
def memory_with_ttl(mock_mcal):
    """Create a MCALMemory with TTL enabled."""
    from mcal_autogen.memory import MCALMemory
    return MCALMemory(
        mcal=mock_mcal,
        user_id="test_user",
        default_ttl_minutes=0.1,  # 6 seconds for testing
    )


# ============================================================================
# Test: Initialization
# ============================================================================

class TestInit:
    """Test MCALMemory initialization."""
    
    def test_basic_init(self, mock_mcal):
        """Test basic initialization with defaults."""
        from mcal_autogen.memory import MCALMemory
        
        mem = MCALMemory(mock_mcal)
        
        assert mem.name == "mcal_memory"
        assert mem.user_id == "default"
        assert mem._max_results == 10
        assert mem._score_threshold == 0.0
        assert mem._default_ttl_minutes is None
        assert mem._enable_goal_tracking is True
        assert mem._include_decisions is True
    
    def test_custom_init(self, mock_mcal):
        """Test initialization with custom parameters."""
        from mcal_autogen.memory import MCALMemory
        
        mem = MCALMemory(
            mcal=mock_mcal,
            user_id="custom_user",
            name="custom_memory",
            max_results=5,
            score_threshold=0.5,
            default_ttl_minutes=30,
            enable_goal_tracking=False,
            include_decisions=False,
        )
        
        assert mem.name == "custom_memory"
        assert mem.user_id == "custom_user"
        assert mem._max_results == 5
        assert mem._score_threshold == 0.5
        assert mem._default_ttl_minutes == 30
        assert mem._enable_goal_tracking is False
        assert mem._include_decisions is False
    
    def test_mcal_property(self, memory, mock_mcal):
        """Test mcal property returns the MCAL instance."""
        assert memory.mcal is mock_mcal


# ============================================================================
# Test: add() method
# ============================================================================

class TestAdd:
    """Test MCALMemory.add() method."""
    
    @pytest.mark.asyncio
    async def test_add_text_content(self, memory):
        """Test adding text content."""
        content = MockMemoryContent(
            content="Test memory content",
            mime_type="text/plain",
        )
        
        await memory.add(content)
        
        assert memory.get_item_count() == 1
    
    @pytest.mark.asyncio
    async def test_add_with_metadata(self, memory):
        """Test adding content with metadata."""
        content = MockMemoryContent(
            content="Decision: Use PostgreSQL",
            mime_type="text/plain",
            metadata={"category": "database", "decision": True},
        )
        
        await memory.add(content)
        
        items = memory.get_all_items()
        assert len(items) == 1
        assert items[0].metadata.get("category") == "database"
    
    @pytest.mark.asyncio
    async def test_add_multiple_items(self, memory):
        """Test adding multiple items."""
        for i in range(5):
            content = MockMemoryContent(content=f"Memory item {i}")
            await memory.add(content)
        
        assert memory.get_item_count() == 5
    
    @pytest.mark.asyncio
    async def test_add_dict_content(self, memory):
        """Test adding dictionary content."""
        content = MockMemoryContent(
            content={"key": "value", "data": [1, 2, 3]},
            mime_type="application/json",
        )
        
        await memory.add(content)
        
        assert memory.get_item_count() == 1
    
    @pytest.mark.asyncio
    async def test_add_cancelled(self, memory):
        """Test that add respects cancellation token."""
        content = MockMemoryContent(content="Should not be added")
        cancel_token = MockCancellationToken(cancelled=True)
        
        await memory.add(content, cancellation_token=cancel_token)
        
        assert memory.get_item_count() == 0
    
    @pytest.mark.asyncio
    async def test_add_with_ttl_metadata(self, memory):
        """Test adding content with TTL in metadata."""
        content = MockMemoryContent(
            content="Temporary content",
            metadata={"ttl_minutes": 0.05},  # 3 seconds
        )
        
        await memory.add(content)
        
        # Should exist immediately
        assert memory.get_item_count() == 1
        
        # Wait for expiration
        time.sleep(4)
        
        # Should be gone after TTL
        assert memory.get_item_count() == 0


# ============================================================================
# Test: query() method
# ============================================================================

class TestQuery:
    """Test MCALMemory.query() method."""
    
    @pytest.mark.asyncio
    async def test_query_empty_memory(self, memory):
        """Test querying empty memory."""
        result = await memory.query("test query")
        
        assert len(result.results) == 0
    
    @pytest.mark.asyncio
    async def test_query_with_match(self, memory, mock_mcal):
        """Test query that finds a match."""
        # Add content through MCAL's add method
        await mock_mcal.add(
            messages=[{"role": "user", "content": "PostgreSQL is great for JSON support"}],
            user_id="test_user",
        )
        
        content = MockMemoryContent(content="PostgreSQL is great for JSON support")
        await memory.add(content)
        
        result = await memory.query("PostgreSQL")
        
        assert len(result.results) > 0
        assert "PostgreSQL" in result.results[0].content
    
    @pytest.mark.asyncio
    async def test_query_with_memory_content(self, memory):
        """Test query using MemoryContent instead of string."""
        content = MockMemoryContent(content="Test content for query")
        await memory.add(content)
        
        query = MockMemoryContent(content="Test content")
        result = await memory.query(query)
        
        # Should match via fallback text search
        assert len(result.results) > 0
    
    @pytest.mark.asyncio
    async def test_query_max_results(self, memory):
        """Test max_results parameter."""
        # Add many items
        for i in range(20):
            content = MockMemoryContent(content=f"Item number {i}")
            await memory.add(content)
        
        result = await memory.query("Item", max_results=5)
        
        assert len(result.results) <= 5
    
    @pytest.mark.asyncio
    async def test_query_cancelled(self, memory):
        """Test that query respects cancellation token."""
        content = MockMemoryContent(content="Test content")
        await memory.add(content)
        
        cancel_token = MockCancellationToken(cancelled=True)
        result = await memory.query("Test", cancellation_token=cancel_token)
        
        assert len(result.results) == 0
    
    @pytest.mark.asyncio
    async def test_query_empty_string(self, memory):
        """Test query with empty string."""
        content = MockMemoryContent(content="Some content")
        await memory.add(content)
        
        result = await memory.query("")
        
        assert len(result.results) == 0


# ============================================================================
# Test: update_context() method
# ============================================================================

class TestUpdateContext:
    """Test MCALMemory.update_context() method."""
    
    @pytest.mark.asyncio
    async def test_update_context_empty_messages(self, memory):
        """Test update_context with no messages."""
        context = MockChatCompletionContext([])
        
        result = await memory.update_context(context)
        
        assert len(result.memories.results) == 0
    
    @pytest.mark.asyncio
    async def test_update_context_with_messages(self, memory):
        """Test update_context adds memories to context."""
        # Add some memories
        content = MockMemoryContent(content="Database choice: PostgreSQL")
        await memory.add(content)
        
        # Create context with a message
        user_msg = MockUserMessage(content="What database should I use?")
        context = MockChatCompletionContext([user_msg])
        
        result = await memory.update_context(context)
        
        # Result should contain the update
        assert isinstance(result.memories, MockMemoryQueryResult)
    
    @pytest.mark.asyncio
    async def test_update_context_adds_system_message(self, memory, mock_mcal):
        """Test that update_context adds a system message."""
        # Add content to MCAL for search
        await mock_mcal.add(
            messages=[{"role": "user", "content": "Use PostgreSQL for the database"}],
            user_id="test_user",
        )
        content = MockMemoryContent(content="Use PostgreSQL for the database")
        await memory.add(content)
        
        user_msg = MockUserMessage(content="PostgreSQL")
        context = MockChatCompletionContext([user_msg])
        
        await memory.update_context(context)
        
        # Context should have additional message if memories found
        messages = await context.get_messages()
        # May have 1 or 2 messages depending on if memories were found
        assert len(messages) >= 1


# ============================================================================
# Test: clear() method
# ============================================================================

class TestClear:
    """Test MCALMemory.clear() method."""
    
    @pytest.mark.asyncio
    async def test_clear_removes_items(self, memory):
        """Test that clear removes all items."""
        for i in range(5):
            content = MockMemoryContent(content=f"Item {i}")
            await memory.add(content)
        
        assert memory.get_item_count() == 5
        
        await memory.clear()
        
        assert memory.get_item_count() == 0
    
    @pytest.mark.asyncio
    async def test_clear_is_user_specific(self, mock_mcal):
        """Test that clear only affects the current user."""
        from mcal_autogen.memory import MCALMemory
        
        mem1 = MCALMemory(mock_mcal, user_id="user1")
        mem2 = MCALMemory(mock_mcal, user_id="user2")
        
        await mem1.add(MockMemoryContent(content="User 1 content"))
        await mem2.add(MockMemoryContent(content="User 2 content"))
        
        await mem1.clear()
        
        # mem1 should be empty, mem2 should still have content
        assert mem1.get_item_count() == 0
        assert mem2.get_item_count() == 1


# ============================================================================
# Test: close() method
# ============================================================================

class TestClose:
    """Test MCALMemory.close() method."""
    
    @pytest.mark.asyncio
    async def test_close_clears_data(self, memory):
        """Test that close clears internal data."""
        content = MockMemoryContent(content="Test content")
        await memory.add(content)
        
        assert memory.get_item_count() == 1
        
        await memory.close()
        
        assert memory.get_item_count() == 0


# ============================================================================
# Test: TTL and Expiration
# ============================================================================

class TestTTL:
    """Test TTL and expiration functionality."""
    
    @pytest.mark.asyncio
    async def test_default_ttl(self, memory_with_ttl):
        """Test that default TTL causes expiration."""
        content = MockMemoryContent(content="Temporary content")
        await memory_with_ttl.add(content)
        
        # Should exist immediately
        assert memory_with_ttl.get_item_count() == 1
        
        # Wait for expiration (6+ seconds for 0.1 minute TTL)
        time.sleep(7)
        
        # Should be expired
        assert memory_with_ttl.get_item_count() == 0
    
    @pytest.mark.asyncio
    async def test_per_item_ttl_override(self, memory):
        """Test per-item TTL override."""
        content = MockMemoryContent(
            content="Short-lived content",
            metadata={"ttl_minutes": 0.05},  # 3 seconds
        )
        await memory.add(content)
        
        # Should exist immediately
        assert memory.get_item_count() == 1
        
        # Wait for expiration
        time.sleep(4)
        
        # Should be expired
        assert memory.get_item_count() == 0
    
    @pytest.mark.asyncio
    async def test_no_ttl_no_expiration(self, memory):
        """Test that items without TTL don't expire."""
        content = MockMemoryContent(content="Permanent content")
        await memory.add(content)
        
        # Should exist
        assert memory.get_item_count() == 1
        
        # Wait a bit
        time.sleep(1)
        
        # Should still exist
        assert memory.get_item_count() == 1


# ============================================================================
# Test: Thread Safety
# ============================================================================

class TestThreadSafety:
    """Test thread safety of operations."""
    
    @pytest.mark.asyncio
    async def test_concurrent_adds(self, memory):
        """Test concurrent add operations."""
        async def add_items(start: int, count: int):
            for i in range(start, start + count):
                content = MockMemoryContent(content=f"Item {i}")
                await memory.add(content)
        
        # Run concurrent adds
        await asyncio.gather(
            add_items(0, 10),
            add_items(10, 10),
            add_items(20, 10),
        )
        
        # All items should be added
        assert memory.get_item_count() == 30


# ============================================================================
# Test: Helper Methods
# ============================================================================

class TestHelperMethods:
    """Test helper/convenience methods."""
    
    @pytest.mark.asyncio
    async def test_add_text(self, memory):
        """Test add_text convenience method."""
        await memory.add_text("Simple text", metadata={"key": "value"})
        
        assert memory.get_item_count() == 1
        items = memory.get_all_items()
        assert items[0].content == "Simple text"
    
    @pytest.mark.asyncio
    async def test_query_text(self, memory):
        """Test query_text convenience method."""
        await memory.add_text("PostgreSQL is recommended")
        await memory.add_text("MySQL is also an option")
        
        results = await memory.query_text("PostgreSQL")
        
        assert isinstance(results, list)
        assert all(isinstance(r, str) for r in results)
    
    def test_get_item_count(self, memory):
        """Test get_item_count method."""
        assert memory.get_item_count() == 0
    
    def test_get_all_items(self, memory):
        """Test get_all_items method."""
        items = memory.get_all_items()
        assert isinstance(items, list)
        assert len(items) == 0


# ============================================================================
# Test: Text Extraction
# ============================================================================

class TestTextExtraction:
    """Test _extract_text method."""
    
    @pytest.mark.asyncio
    async def test_extract_from_string(self, memory):
        """Test extracting text from string."""
        text = memory._extract_text("Hello world")
        assert text == "Hello world"
    
    @pytest.mark.asyncio
    async def test_extract_from_memory_content(self, memory):
        """Test extracting text from MemoryContent."""
        content = MockMemoryContent(content="Content text")
        text = memory._extract_text(content)
        assert text == "Content text"
    
    @pytest.mark.asyncio
    async def test_extract_from_dict_content(self, memory):
        """Test extracting text from dict content."""
        content = MockMemoryContent(content={"text": "Dict text"})
        text = memory._extract_text(content)
        assert text == "Dict text"


# ============================================================================
# Test: Issue #102 - API alignment fixes
# ============================================================================

class TestIssue102APIAlignment:
    """Tests for Issue #102: Verify correct MCAL API usage."""
    
    @pytest.mark.asyncio
    async def test_add_calls_mcal_add_with_messages(self, memory):
        """Verify add() calls MCAL.add(messages=..., user_id=...)."""
        memory._mcal.add = AsyncMock(return_value=MagicMock())
        
        content = MockMemoryContent(content="Test content for MCAL")
        await memory.add(content)
        
        # Should call await self._mcal.add(messages=[...], user_id=...)
        memory._mcal.add.assert_awaited_once()
        call_kwargs = memory._mcal.add.call_args.kwargs
        assert "messages" in call_kwargs
        assert "user_id" in call_kwargs
        assert call_kwargs["messages"][0]["role"] == "user"
        assert call_kwargs["messages"][0]["content"] == "Test content for MCAL"
        assert call_kwargs["user_id"] == "test_user"
    
    @pytest.mark.asyncio
    async def test_query_calls_mcal_search_with_limit(self, memory):
        """Verify query() calls MCAL.search(limit=...) not top_k=."""
        mock_entry = MagicMock()
        mock_entry.content = "result"
        mock_entry.score = 0.9
        mock_entry.metadata = {}
        
        mock_result = MagicMock()
        mock_result.memories = [mock_entry]
        memory._mcal.search = AsyncMock(return_value=mock_result)
        
        await memory.query("test query", max_results=5)
        
        memory._mcal.search.assert_awaited_once()
        call_kwargs = memory._mcal.search.call_args.kwargs
        assert "limit" in call_kwargs, "Should use limit= parameter"
        assert call_kwargs["limit"] == 5
        assert "top_k" not in call_kwargs, "Should NOT use top_k="
    
    @pytest.mark.asyncio
    async def test_query_awaits_mcal_search(self, memory, mock_mcal):
        """Verify query() properly awaits the async search call."""
        # Add a memory through the mock
        await mock_mcal.add(
            messages=[{"role": "user", "content": "async test data"}],
            user_id="test_user",
        )
        
        # This should work without errors (would fail without await)
        result = await memory.query("async test")
        assert len(result.results) > 0
    
    @pytest.mark.asyncio
    async def test_query_reads_memories_not_raw_result(self, memory):
        """Verify query() accesses .memories on SearchResult."""
        mock_entry = MagicMock()
        mock_entry.content = "found content"
        mock_entry.score = 0.85
        mock_entry.metadata = {"key": "val"}
        
        mock_result = MagicMock()
        mock_result.memories = [mock_entry]
        # Ensure iterating the result directly would fail
        mock_result.__iter__ = MagicMock(side_effect=TypeError("Not iterable"))
        memory._mcal.search = AsyncMock(return_value=mock_result)
        
        result = await memory.query("test")
        
        assert len(result.results) == 1
        assert result.results[0].content == "found content"
    
    @pytest.mark.asyncio
    async def test_add_awaits_mcal_add(self, memory):
        """Verify add() properly awaits the async MCAL.add() call."""
        call_log = []
        
        async def mock_add(**kwargs):
            call_log.append("awaited")
            return MagicMock()
        
        memory._mcal.add = mock_add
        
        content = MockMemoryContent(content="await test")
        await memory.add(content)
        
        assert call_log == ["awaited"], "MCAL.add() should be awaited"


# ============================================================================
# Test: Error Handling
# ============================================================================

class TestErrorHandling:
    """Test error handling scenarios."""
    
    @pytest.mark.asyncio
    async def test_mcal_add_failure_graceful(self, memory):
        """Test that MCAL add failure doesn't crash."""
        # Make MCAL raise an exception
        memory._mcal.add = AsyncMock(side_effect=Exception("MCAL error"))
        
        content = MockMemoryContent(content="Test content")
        
        # Should not raise, but silently continue
        await memory.add(content)
        
        # In-memory storage should still work
        assert memory.get_item_count() == 1
    
    @pytest.mark.asyncio
    async def test_mcal_search_failure_fallback(self, memory):
        """Test that MCAL search failure falls back to in-memory."""
        # Add content
        content = MockMemoryContent(content="Fallback test content")
        await memory.add(content)
        
        # Make MCAL search raise an exception
        memory._mcal.search = AsyncMock(side_effect=Exception("Search error"))
        
        # Should fall back to in-memory search
        result = await memory.query("Fallback")
        
        # Should still find the content via fallback
        assert len(result.results) > 0


# ============================================================================
# Run tests
# ============================================================================

if __name__ == "__main__":
    pytest.main([__file__, "-v"])
