"""
Tests for Core FinancialSituationMemory Class

This file tests the core FinancialSituationMemory class methods and behavior.

NOT covered here (see other test files):
- Ticker-specific isolation → test_memory_isolation.py
- Ticker sanitization → test_memory_isolation.py
- create_memory_instances() → test_memory_isolation.py
- cleanup_all_memories() → test_memory_isolation.py
- get_all_memory_stats() → test_memory_isolation.py
- Graph integration → test_memory_integration.py

Focus:
- Memory initialization and availability detection
- Situation storage and retrieval
- Memory cleanup (per-instance)
- Statistics (per-instance)
"""

import socket
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from src.llms import _LazyRateLimiterProxy
from src.memory import (
    FinancialSituationMemory,
    cleanup_all_memories,
    get_all_memory_stats,
    get_ticker_memory_stats,
)


@pytest.fixture(autouse=True)
def reset_shared_memory_state():
    FinancialSituationMemory._reset_shared_state_for_tests()
    yield
    FinancialSituationMemory._reset_shared_state_for_tests()


class TestFinancialSituationMemoryInitialization:
    """Test memory initialization under various conditions."""

    def test_init_without_api_key(self):
        """Memory should gracefully handle missing API key.

        Note: We mock config.get_google_api_key() to return empty string,
        which simulates the case where GOOGLE_API_KEY is not set.
        This is cleaner than patching os.environ and reloading modules.
        """
        with patch("src.memory.config") as mock_config:
            mock_config.get_google_api_key.return_value = ""

            memory = FinancialSituationMemory("test_memory")

            assert memory.name == "test_memory"
            assert not memory.available
            assert memory.embeddings is None
            assert memory.situation_collection is None

    def test_init_with_embedding_failure(self):
        """Memory should handle embedding initialization failures."""
        with patch("src.memory.config") as mock_config:
            mock_config.get_google_api_key.return_value = "test-key"
            with patch(
                "src.memory.GoogleGenerativeAIEmbeddings",
                side_effect=Exception("API error"),
            ):
                memory = FinancialSituationMemory("test_memory")

                assert not memory.available
                assert memory.embeddings is None

    def test_init_with_chromadb_failure(self):
        """Memory should handle ChromaDB connection failures."""
        with patch("src.memory.config") as mock_config:
            mock_config.get_google_api_key.return_value = "test-key"
            with patch(
                "src.memory.GoogleGenerativeAIEmbeddings"
            ) as mock_embeddings_class:
                mock_embeddings = MagicMock()
                mock_embeddings.embed_query.return_value = [0.1] * 768
                mock_embeddings_class.return_value = mock_embeddings

                with patch(
                    "chromadb.PersistentClient",
                    side_effect=Exception("Connection failed"),
                ):
                    memory = FinancialSituationMemory("test_memory")

                    assert not memory.available
                    assert memory.situation_collection is None

    def test_successful_init(self):
        """Memory should initialize successfully with valid configuration."""
        FinancialSituationMemory._reset_shared_state_for_tests()
        # Mock config getter directly (os.environ patching no longer works due to SecretStr)
        with patch("src.memory.config") as mock_config:
            mock_config.get_google_api_key.return_value = "test-key"

            with patch(
                "src.memory.GoogleGenerativeAIEmbeddings"
            ) as mock_embeddings_class:
                mock_embeddings = MagicMock()
                mock_embeddings.embed_query.return_value = [0.1] * 768
                mock_embeddings_class.return_value = mock_embeddings

                with patch("chromadb.PersistentClient") as mock_client_class:
                    mock_client = MagicMock()
                    mock_collection = MagicMock()
                    mock_collection.count.return_value = 0
                    mock_client.get_or_create_collection.return_value = mock_collection
                    mock_client_class.return_value = mock_client

                    memory = FinancialSituationMemory("test_memory")

                    assert memory.name == "test_memory"
                    assert memory.available
                    assert memory.embeddings is not None
                    assert memory.situation_collection is not None
        FinancialSituationMemory._reset_shared_state_for_tests()

    def test_reuses_shared_embeddings_and_chroma_client(self):
        """Multiple memory instances should share heavy embedding/client resources."""
        FinancialSituationMemory._reset_shared_state_for_tests()
        with patch("src.memory.config") as mock_config:
            mock_config.get_google_api_key.return_value = "test-key"
            mock_config.chroma_persist_directory = "test-chroma"

            with patch(
                "src.memory.GoogleGenerativeAIEmbeddings"
            ) as mock_embeddings_class:
                mock_embeddings = MagicMock()
                mock_embeddings.embed_query.return_value = [0.1] * 768
                mock_embeddings_class.return_value = mock_embeddings

                with patch("chromadb.PersistentClient") as mock_client_class:
                    mock_client = MagicMock()
                    mock_collection_a = MagicMock()
                    mock_collection_a.count.return_value = 0
                    mock_collection_b = MagicMock()
                    mock_collection_b.count.return_value = 0
                    mock_client.list_collections.return_value = []
                    mock_client.get_or_create_collection.side_effect = [
                        mock_collection_a,
                        mock_collection_b,
                    ]
                    mock_client_class.return_value = mock_client

                    memory_a = FinancialSituationMemory("test_memory_a")
                    memory_b = FinancialSituationMemory("test_memory_b")

                    assert memory_a.available is True
                    assert memory_b.available is True
                    assert mock_embeddings_class.call_count == 1
                    assert mock_client_class.call_count == 1
        FinancialSituationMemory._reset_shared_state_for_tests()

    def test_init_with_embedding_healthcheck_failure_is_degraded(self):
        """Memory should not report available when embedding healthcheck fails."""
        with patch("src.memory.config") as mock_config:
            mock_config.get_google_api_key.return_value = "test-key"

            with patch(
                "src.memory.GoogleGenerativeAIEmbeddings"
            ) as mock_embeddings_class:
                mock_embeddings = MagicMock()
                mock_embeddings.embed_query.side_effect = socket.gaierror(
                    8, "nodename nor servname provided, or not known"
                )
                mock_embeddings_class.return_value = mock_embeddings

                with patch("chromadb.PersistentClient") as mock_client_class:
                    mock_client = MagicMock()
                    mock_collection = MagicMock()
                    mock_collection.count.return_value = 0
                    mock_client.get_or_create_collection.return_value = mock_collection
                    mock_client_class.return_value = mock_client

                    memory = FinancialSituationMemory("test_memory")

                    assert memory.chroma_available is True
                    assert memory.embeddings_available is False
                    assert memory.available is False

    def test_get_ticker_memory_stats_reads_existing_collections_without_recreating(
        self,
    ):
        """The read-only stats helper should inspect existing collections only."""
        with patch("src.memory.config") as mock_config:
            mock_config.chroma_persist_directory = "test-chroma"

            mock_client = MagicMock()
            existing = MagicMock()
            existing.name = "1308_HK_bull_memory"
            mock_collection = MagicMock()
            mock_collection.count.return_value = 3
            mock_client.list_collections.return_value = [existing]
            mock_client.get_collection.return_value = mock_collection

            with patch.object(
                FinancialSituationMemory,
                "_get_shared_chroma_client",
                return_value=mock_client,
            ) as mock_get_client:
                stats = get_ticker_memory_stats("1308.HK")

        mock_get_client.assert_called_once()
        assert stats["bull_researcher"]["available"] is True
        assert stats["bull_researcher"]["count"] == 3
        assert stats["bear_researcher"]["available"] is False


class TestSituationStorage:
    """Test adding financial situations to memory."""

    @pytest.mark.asyncio
    async def test_add_situations_unavailable(self):
        """add_situations should return False when memory unavailable."""
        memory = FinancialSituationMemory("test_memory")
        memory.available = False

        result = await memory.add_situations(["Test situation"])

        assert not result

    @pytest.mark.asyncio
    async def test_add_situations_empty_list(self):
        """add_situations should return False for empty list."""
        memory = FinancialSituationMemory("test_memory")
        memory.available = True

        result = await memory.add_situations([])

        assert not result

    @pytest.mark.asyncio
    async def test_add_situations_success(self):
        """add_situations should successfully store situations."""
        memory = FinancialSituationMemory("test_memory")
        memory.available = True
        memory._get_embedding = AsyncMock(return_value=[0.1] * 768)
        memory.situation_collection = MagicMock()

        situations = ["AAPL strong buy signal", "Market momentum positive"]
        result = await memory.add_situations(situations)

        assert result

        # Verify collection.add was called
        assert memory.situation_collection.add.called
        call_kwargs = memory.situation_collection.add.call_args[1]

        # Check that all required fields were provided
        assert "embeddings" in call_kwargs
        assert "documents" in call_kwargs
        assert "ids" in call_kwargs
        assert "metadatas" in call_kwargs

        # Check correct number of items
        assert len(call_kwargs["documents"]) == 2
        assert len(call_kwargs["embeddings"]) == 2

    @pytest.mark.asyncio
    async def test_get_embedding_logs_rate_limiter_fallback(self):
        """Rate-limiter failures should be visible before direct embedding fallback."""

        class BrokenLimiter:
            async def __aenter__(self):
                raise RuntimeError("limiter unavailable")

            async def __aexit__(self, exc_type, exc, tb):
                return False

        memory = FinancialSituationMemory("test_memory")
        memory.available = True
        memory.embeddings = MagicMock()
        memory.embeddings.aembed_query = AsyncMock(return_value=[0.1] * 768)

        with patch("src.llms.GLOBAL_RATE_LIMITER", BrokenLimiter()):
            with patch("src.memory.logger") as mock_logger:
                embedding = await memory._get_embedding("hello world")

        assert embedding == [0.1] * 768
        mock_logger.warning.assert_called_once()
        kwargs = mock_logger.warning.call_args.kwargs
        assert kwargs["fallback"] == "direct_embed_query"
        assert kwargs["failure_kind"] in {
            "application_error",
            "connect_error",
            "timeout",
            "unknown_provider_error",
        }
        assert kwargs["error_type"] == "RuntimeError"
        assert kwargs["exc_info"] is True

    @pytest.mark.asyncio
    async def test_get_embedding_accepts_lazy_rate_limiter_proxy(self):
        """The shared lazy proxy should behave as an async context manager."""

        class TrackingLimiter:
            def __init__(self):
                self.acquired = 0

            async def aacquire(self, *, blocking=True):
                self.acquired += 1
                return True

        limiter = TrackingLimiter()
        memory = FinancialSituationMemory("test_memory")
        memory.available = True
        memory.embeddings = MagicMock()
        memory.embeddings.aembed_query = AsyncMock(return_value=[0.2] * 768)

        with patch(
            "src.llms.GLOBAL_RATE_LIMITER", _LazyRateLimiterProxy(lambda: limiter)
        ):
            with patch("src.memory.logger") as mock_logger:
                embedding = await memory._get_embedding("hello world")

        assert embedding == [0.2] * 768
        assert limiter.acquired == 1
        mock_logger.warning.assert_not_called()


class TestSituationQuerying:
    """Test querying similar situations from memory."""

    @pytest.mark.asyncio
    async def test_query_unavailable(self):
        """query_similar_situations should return empty list when unavailable."""
        memory = FinancialSituationMemory("test_memory")
        memory.available = False

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

        assert results == []

    @pytest.mark.asyncio
    async def test_query_success(self):
        """query_similar_situations should return formatted results."""
        memory = FinancialSituationMemory("test_memory")
        memory.available = True
        memory._get_embedding = AsyncMock(return_value=[0.1] * 768)
        memory.situation_collection = MagicMock()

        # Mock ChromaDB query response
        memory.situation_collection.query.return_value = {
            "documents": [["AAPL strong fundamentals", "MSFT growing revenue"]],
            "metadatas": [[{"ticker": "AAPL"}, {"ticker": "MSFT"}]],
            "distances": [[0.1, 0.2]],
        }

        results = await memory.query_similar_situations("tech stocks", n_results=2)

        assert len(results) == 2
        assert results[0]["document"] == "AAPL strong fundamentals"
        assert results[0]["metadata"]["ticker"] == "AAPL"
        assert results[0]["distance"] == 0.1
        assert results[1]["document"] == "MSFT growing revenue"

    @pytest.mark.parametrize(
        "message",
        [
            "Error getting collection: Collection [foo] does not exist.",
            "Collection not found",
        ],
    )
    @pytest.mark.asyncio
    async def test_query_missing_collection_does_not_log_error(self, message):
        """Missing collections should be treated as an expected empty-slate state."""
        memory = FinancialSituationMemory("test_memory")
        memory.available = True
        memory._get_embedding = AsyncMock(return_value=[0.1] * 768)
        memory.situation_collection = MagicMock()
        memory.situation_collection.query.side_effect = Exception(message)

        with (
            patch("src.memory.logger") as mock_logger,
            patch("src.memory._record_capture_memory_event") as mock_capture_event,
        ):
            results = await memory.query_similar_situations("tech stocks", n_results=2)

        assert results == []
        mock_logger.debug.assert_called_once()
        assert (
            mock_logger.debug.call_args.args[0]
            == "query_similar_situations_no_collection"
        )
        mock_logger.error.assert_not_called()
        payload = mock_capture_event.call_args.args[0]
        assert payload["event"] == "query_similar_situations_failed"
        assert payload["error_kind"] == "missing_collection"
        assert payload["available"] is True

    @pytest.mark.asyncio
    async def test_query_non_missing_failure_logs_error(self):
        """Unexpected query failures should still log an error."""
        memory = FinancialSituationMemory("test_memory")
        memory.available = True
        memory._get_embedding = AsyncMock(return_value=[0.1] * 768)
        memory.situation_collection = MagicMock()
        memory.situation_collection.query.side_effect = Exception("connection reset")

        with (
            patch("src.memory.logger") as mock_logger,
            patch("src.memory._record_capture_memory_event") as mock_capture_event,
        ):
            results = await memory.query_similar_situations("tech stocks")

        assert results == []
        mock_logger.error.assert_called_once()
        assert mock_logger.error.call_args.args[0] == "query_similar_situations_failed"
        payload = mock_capture_event.call_args.args[0]
        assert payload["error_kind"] == "query_failure"
        assert payload["available"] is True


class TestGetRelevantMemory:
    """Test high-level memory retrieval for agent context."""

    @pytest.mark.asyncio
    async def test_get_relevant_memory_no_results(self):
        """get_relevant_memory should handle no results gracefully."""
        memory = FinancialSituationMemory("test_memory")
        memory.available = True
        memory.query_similar_situations = AsyncMock(return_value=[])

        result = await memory.get_relevant_memory("AAPL", "analysis")

        assert "No relevant past memories found" in result

    @pytest.mark.asyncio
    async def test_get_relevant_memory_success(self):
        """get_relevant_memory should format results for display."""
        memory = FinancialSituationMemory("test_memory")
        memory.available = True

        mock_results = [
            {
                "document": "AAPL shows strong fundamentals with growing revenue",
                "metadata": {"ticker": "AAPL", "timestamp": "2024-01-01T00:00:00"},
                "distance": 0.1,
            }
        ]
        memory.query_similar_situations = AsyncMock(return_value=mock_results)

        result = await memory.get_relevant_memory("AAPL", "fundamental analysis")

        assert "Relevant past memories for AAPL" in result
        assert "AAPL shows strong fundamentals" in result


class TestMemoryCleanup:
    """
    Test memory cleanup functionality.
    Note: cleanup_all_memories returns a dict[str, int], NOT an int.
    """

    def test_cleanup_unavailable(self):
        """clear_old_memories (instance method) should return dict when unavailable (it re-initializes client)."""
        memory = FinancialSituationMemory("test_memory")
        memory.available = False

        # Instance method now returns dict[str, int]
        # Even if available=False, it creates a new client to perform cleanup
        results = memory.clear_old_memories(days_to_keep=30)
        assert isinstance(results, dict)

    @patch("chromadb.PersistentClient")
    def test_cleanup_no_old_memories(self, mock_client_cls):
        """cleanup_all_memories should return dictionary of zeros when no old memories."""
        mock_client = MagicMock()
        mock_collection = MagicMock()
        mock_collection.name = "test_memory"

        # FIXED: Explicitly set count return value to avoid MagicMock comparison error
        mock_collection.count.return_value = 0

        # Mock list_collections returning the collection
        mock_client.list_collections.return_value = [mock_collection]

        # Mock get() returning empty
        mock_collection.get.return_value = {"ids": [], "metadatas": []}

        mock_client_cls.return_value = mock_client

        results = cleanup_all_memories(days=0)

        # Should return a dict mapping collection names to counts
        assert isinstance(results, dict)
        # Verify the mocked collection name is in the results
        assert results.get("test_memory") == 0

    @patch("chromadb.PersistentClient")
    def test_cleanup_with_old_memories(self, mock_client_cls):
        """cleanup_all_memories should return correct counts in dictionary."""
        mock_client = MagicMock()
        mock_collection = MagicMock()
        mock_collection.name = "test_memory"

        # When count() is called on the mock, return 5
        mock_collection.count.return_value = 5

        # Mock list_collections
        mock_client.list_collections.return_value = [mock_collection]

        mock_client_cls.return_value = mock_client

        # Run cleanup with days=0 (should delete everything)
        results = cleanup_all_memories(days=0)

        assert isinstance(results, dict)
        # Check that we got 5 deletions for test_memory
        # The previous failure asserted a dict != 1, so we must check the dict value
        assert results.get("test_memory") == 5
        mock_client.delete_collection.assert_called_with("test_memory")


class TestMemoryStats:
    """Test memory statistics retrieval."""

    def test_stats_unavailable(self):
        """get_stats should return unavailable status."""
        # Mock config to simulate missing API key
        with patch("src.memory.config") as mock_config:
            mock_config.get_google_api_key.return_value = ""

            memory = FinancialSituationMemory("test_memory")

            # Double check initialization state
            assert not memory.available

            stats = memory.get_stats()

            assert not stats["available"]
            assert stats["name"] == "test_memory"
            assert stats["count"] == 0

    def test_stats_success(self):
        """get_stats should return collection information."""
        memory = FinancialSituationMemory("test_memory")
        memory.available = True
        memory.situation_collection = MagicMock()
        memory.situation_collection.count.return_value = 42

        stats = memory.get_stats()

        assert stats["available"]
        assert stats["name"] == "test_memory"
        assert stats["count"] == 42

    @pytest.mark.parametrize(
        "message",
        [
            "Error getting collection: Collection [foo] does not exist.",
            "Collection not found",
        ],
    )
    def test_stats_missing_collection_variants_return_deleted(self, message):
        """Missing collections should be reported as deleted, not fatal."""
        memory = FinancialSituationMemory("test_memory")
        memory.available = True
        memory.situation_collection = MagicMock()
        memory.situation_collection.count.side_effect = Exception(message)

        with patch("src.memory.logger") as mock_logger:
            stats = memory.get_stats()

        assert stats == {
            "available": False,
            "name": "test_memory",
            "count": 0,
            "status": "deleted",
        }
        mock_logger.debug.assert_called_once_with(
            "collection_deleted_externally", collection="test_memory"
        )
        mock_logger.error.assert_not_called()


class TestAllMemoryStats:
    """Test global memory stats collection edge cases."""

    @pytest.mark.parametrize(
        "message",
        [
            "Error getting collection: Collection [foo] does not exist.",
            "Collection not found",
        ],
    )
    @patch("chromadb.PersistentClient")
    def test_get_all_stats_skips_missing_collection_variants(
        self, mock_client_cls, message
    ):
        """Missing collections should be skipped regardless of Chroma message wording."""
        mock_client = MagicMock()
        mock_client.list_collections.return_value = ["missing_collection"]
        mock_client.get_collection.side_effect = Exception(message)
        mock_client_cls.return_value = mock_client

        stats = get_all_memory_stats()

        assert stats == {}


class TestAddSituationsCollectionNotFoundRetry:
    """May 2026 2099.HK incident: a stale `situation_collection` handle
    (e.g. across a model-mismatch deletion in a concurrent process)
    surfaces as `add_situations_failed` (NotFoundError) and a downstream
    `rejection_record_storage_failed` warning. The retry recovers from
    that by re-fetching the collection via `get_or_create_collection`."""

    @pytest.mark.asyncio
    async def test_retries_on_collection_notfound_and_recovers(self):
        from src.memory import FinancialSituationMemory

        memory = FinancialSituationMemory.__new__(FinancialSituationMemory)
        memory.name = "lessons_learned"
        memory.available = True
        memory._EMBEDDING_MODEL = "test-model"
        memory._EMBEDDING_DIMENSION = 768

        # First add() raises a NotFoundError-like exception; second succeeds.
        first_collection = MagicMock()
        first_collection.add.side_effect = Exception(
            "Error getting collection: Collection [lessons_learned] does not exist."
        )
        recovered_collection = MagicMock()
        recovered_collection.add.return_value = None

        memory.situation_collection = first_collection
        memory.chroma_client = MagicMock()
        memory.chroma_client.get_or_create_collection.return_value = (
            recovered_collection
        )

        async def fake_embedding(_text):
            return [0.1] * 768

        memory._get_embedding = fake_embedding

        # Mock inspection to allow content through.
        with patch(
            "src.memory.get_current_inspection_service"
        ) as mock_inspection_service:
            mock_inspection = MagicMock()
            mock_inspection.check = AsyncMock(return_value="approved content")
            mock_inspection_service.return_value = mock_inspection

            ok = await memory.add_situations(["test situation"], [{"ticker": "T"}])

        assert ok is True
        # Both the original and the recovered collection saw an .add() call.
        assert first_collection.add.call_count == 1
        assert recovered_collection.add.call_count == 1
        memory.chroma_client.get_or_create_collection.assert_called_once()

    @pytest.mark.asyncio
    async def test_propagates_non_notfound_errors_without_retry(self):
        """Other errors (e.g. validation) must NOT trigger a refetch retry —
        we'd just be papering over real bugs. Caller sees the standard
        `add_situations_failed` log + False return value."""
        from src.memory import FinancialSituationMemory

        memory = FinancialSituationMemory.__new__(FinancialSituationMemory)
        memory.name = "lessons_learned"
        memory.available = True
        memory._EMBEDDING_MODEL = "test-model"
        memory._EMBEDDING_DIMENSION = 768

        collection = MagicMock()
        collection.add.side_effect = ValueError("schema mismatch")
        memory.situation_collection = collection
        memory.chroma_client = MagicMock()

        async def fake_embedding(_text):
            return [0.1] * 768

        memory._get_embedding = fake_embedding

        with patch(
            "src.memory.get_current_inspection_service"
        ) as mock_inspection_service:
            mock_inspection = MagicMock()
            mock_inspection.check = AsyncMock(return_value="approved")
            mock_inspection_service.return_value = mock_inspection

            ok = await memory.add_situations(["test"], [{}])

        assert ok is False
        assert collection.add.call_count == 1
        memory.chroma_client.get_or_create_collection.assert_not_called()
