"""
Tests for embedding dimension cache bug (Issue #412).

Tests verify that embedding_dimension is properly restored when retrieving
cached models from _MODEL_CACHE across all three embedding types:
- External embedding API
- ONNX embedding
- SentenceTransformer embedding
"""

import pytest
import pytest_asyncio
import tempfile
import os
import shutil
from unittest.mock import Mock, patch, MagicMock

# Skip tests if sqlite-vec is not available
try:
    import sqlite_vec
    SQLITE_VEC_AVAILABLE = True
except ImportError:
    SQLITE_VEC_AVAILABLE = False

if SQLITE_VEC_AVAILABLE:
    from src.mcp_memory_service.storage.sqlite_vec import (
        SqliteVecMemoryStorage,
        _MODEL_CACHE,
        _DIMENSION_CACHE,
        _EMBEDDING_CACHE
    )

pytestmark = pytest.mark.skipif(not SQLITE_VEC_AVAILABLE, reason="sqlite-vec not available")


class TestEmbeddingDimensionCache:
    """Test suite for embedding dimension cache restoration."""

    @pytest_asyncio.fixture
    async def temp_db_path(self):
        """Create a temporary database path."""
        temp_dir = tempfile.mkdtemp()
        db_path = os.path.join(temp_dir, "test_dimension_cache.db")
        yield db_path
        shutil.rmtree(temp_dir, ignore_errors=True)

    @pytest_asyncio.fixture(autouse=True)
    async def clear_caches(self):
        """Clear global caches before each test."""
        _MODEL_CACHE.clear()
        _DIMENSION_CACHE.clear()
        _EMBEDDING_CACHE.clear()
        yield
        _MODEL_CACHE.clear()
        _DIMENSION_CACHE.clear()
        _EMBEDDING_CACHE.clear()

    @pytest.mark.asyncio
    async def test_external_embedding_dimension_restored_from_cache(self, temp_db_path):
        """
        Test that embedding_dimension is restored when retrieving external
        embedding model from cache.

        RED phase: This test should FAIL because the dimension is not restored.
        """
        # Mock external embedding model with 768 dimensions
        mock_model = MagicMock()
        mock_model.embedding_dimension = 768
        mock_model.encode_batch = MagicMock(return_value=[[0.1] * 768])

        with patch.dict(os.environ, {
            'MCP_EXTERNAL_EMBEDDING_URL': 'http://localhost:8890/v1/embeddings',
            'MCP_EXTERNAL_EMBEDDING_MODEL': 'test-model-768',
            'MCP_MEMORY_STORAGE_BACKEND': 'sqlite_vec'
        }):
            with patch('src.mcp_memory_service.embeddings.external_api.get_external_embedding_model', return_value=mock_model):
                # First initialization - loads and caches model with 768 dimensions
                storage1 = SqliteVecMemoryStorage(temp_db_path)
                await storage1.initialize()

                assert storage1.embedding_dimension == 768, \
                    f"First init: expected dimension 768, got {storage1.embedding_dimension}"

                # Close first storage
                if storage1.conn:
                    storage1.conn.close()

                # Second initialization - retrieves from cache
                # BUG: dimension should be 768 but will be 384 (default)
                storage2 = SqliteVecMemoryStorage(temp_db_path)
                await storage2.initialize()

                # This assertion SHOULD FAIL in RED phase
                assert storage2.embedding_dimension == 768, \
                    f"Second init (cached): expected dimension 768, got {storage2.embedding_dimension}"

                # Close second storage
                if storage2.conn:
                    storage2.conn.close()

    @pytest.mark.asyncio
    async def test_onnx_embedding_dimension_restored_from_cache(self, temp_db_path):
        """
        Test that embedding_dimension is restored when retrieving ONNX
        embedding model from cache.

        RED phase: This test should FAIL because the dimension is not restored.
        """
        # Mock ONNX model with non-default dimension (512)
        mock_onnx_model = MagicMock()
        mock_onnx_model.embedding_dimension = 512
        mock_onnx_model.encode_batch = MagicMock(return_value=[[0.1] * 512])

        with patch.dict(os.environ, {
            'MCP_MEMORY_USE_ONNX': 'true',
            'MCP_MEMORY_STORAGE_BACKEND': 'sqlite_vec'
        }):
            with patch('src.mcp_memory_service.embeddings.get_onnx_embedding_model', return_value=mock_onnx_model):
                # First initialization - loads and caches ONNX model with 512 dimensions
                storage1 = SqliteVecMemoryStorage(temp_db_path)
                await storage1.initialize()

                assert storage1.embedding_dimension == 512, \
                    f"First init: expected dimension 512, got {storage1.embedding_dimension}"

                # Close first storage
                if storage1.conn:
                    storage1.conn.close()

                # Second initialization - retrieves from cache
                # BUG: dimension should be 512 but will be 384 (default)
                storage2 = SqliteVecMemoryStorage(temp_db_path)
                await storage2.initialize()

                # This assertion SHOULD FAIL in RED phase
                assert storage2.embedding_dimension == 512, \
                    f"Second init (cached): expected dimension 512, got {storage2.embedding_dimension}"

                # Close second storage
                if storage2.conn:
                    storage2.conn.close()

    @pytest.mark.asyncio
    async def test_sentence_transformer_dimension_restored_from_cache(self, temp_db_path):
        """
        Test that embedding_dimension is restored when retrieving SentenceTransformer
        model from cache using the actual default model (all-MiniLM-L6-v2, 384 dims).

        This test verifies cache behavior works correctly for the default case.
        """
        with patch.dict(os.environ, {
            'MCP_MEMORY_USE_ONNX': 'false',
            'MCP_MEMORY_STORAGE_BACKEND': 'sqlite_vec'
        }):
            # First initialization - loads and caches default model (384 dimensions)
            storage1 = SqliteVecMemoryStorage(temp_db_path)
            await storage1.initialize()

            first_dimension = storage1.embedding_dimension
            assert first_dimension == 384, \
                f"First init: expected default dimension 384, got {first_dimension}"

            # Close first storage
            if storage1.conn:
                storage1.conn.close()

            # Second initialization - retrieves from cache
            # Should restore the same dimension (384)
            storage2 = SqliteVecMemoryStorage(temp_db_path)
            await storage2.initialize()

            # This should pass if dimension caching works
            assert storage2.embedding_dimension == first_dimension, \
                f"Second init (cached): expected dimension {first_dimension}, got {storage2.embedding_dimension}"

            # Close second storage
            if storage2.conn:
                storage2.conn.close()

    @pytest.mark.asyncio
    async def test_dimension_cache_handles_missing_dimension(self, temp_db_path):
        """
        Test graceful handling when dimension is missing from cache.

        This ensures backward compatibility if cache is partially populated.
        """
        # This test verifies the fix handles edge cases properly
        # Will be written after GREEN phase to ensure proper behavior
        pass
