Session Buddy

"""Tests for HNSW vector indexing performance optimization.

Tests HNSW (Hierarchical Navigable Small World) index creation, usage,
and performance improvements for vector similarity search.
"""

from __future__ import annotations

import typing as t
from datetime import UTC, datetime
from pathlib import Path

import pytest

from session_buddy.adapters.reflection_adapter_oneiric import (
    ReflectionDatabaseAdapterOneiric,
)
from session_buddy.adapters.settings import ReflectionAdapterSettings


class TestHNSWIndexCreation:
    """Test HNSW index creation and configuration."""

    @pytest.fixture
    def temp_db_path(self, tmp_path: Path) -> Path:
        """Create temporary database path."""
        return tmp_path / "test_reflection.duckdb"

    @pytest.fixture
    def adapter_settings(self, temp_db_path: Path) -> ReflectionAdapterSettings:
        """Create adapter settings with HNSW enabled."""
        return ReflectionAdapterSettings(
            database_path=temp_db_path,
            collection_name="test_hnsw",
            enable_hnsw_index=True,
            hnsw_m=16,
            hnsw_ef_construction=200,
            hnsw_ef_search=64,
        )

    @pytest.mark.asyncio
    async def test_hnsw_index_creation_on_init(self, adapter_settings: ReflectionAdapterSettings) -> None:
        """Test that HNSW indexes are created during adapter initialization."""
        adapter = ReflectionDatabaseAdapterOneiric(settings=adapter_settings)

        # Initialize should create tables and HNSW indexes
        await adapter.initialize()

        # Check that HNSW indexes exist
        result = adapter.conn.execute(
            "SELECT * FROM duckdb_indexes() WHERE index_name LIKE '%hnsw%'"
        ).fetchall()

        # Should have 2 HNSW indexes (conversations and reflections)
        # Column 4 is index_name (verified from duckdb_indexes() schema)
        hnsw_indexes = [row for row in result if "hnsw" in row[4].lower()]

        assert len(hnsw_indexes) >= 2, f"Expected at least 2 HNSW indexes, got {len(hnsw_indexes)}"

        await adapter.aclose()

    @pytest.mark.asyncio
    async def test_hnsw_index_disabled_when_setting_false(self, temp_db_path: Path) -> None:
        """Test that HNSW indexes are not created when enable_hnsw_index=False."""
        settings = ReflectionAdapterSettings(
            database_path=temp_db_path,
            collection_name="test_no_hnsw",
            enable_hnsw_index=False,
        )

        adapter = ReflectionDatabaseAdapterOneiric(settings=settings)
        await adapter.initialize()

        # Check that HNSW indexes do NOT exist
        # Filter for indexes ending with '_hnsw' to avoid false matches
        result = adapter.conn.execute(
            "SELECT index_name FROM duckdb_indexes() WHERE index_name LIKE '%_hnsw'"
        ).fetchall()

        # Should have no HNSW indexes when disabled
        hnsw_indexes = [row[0] for row in result]
        assert len(hnsw_indexes) == 0, f"Expected no HNSW indexes when disabled, got {len(hnsw_indexes)}: {hnsw_indexes}"

        await adapter.aclose()

    @pytest.mark.asyncio
    async def test_hnsw_index_parameters_respected(self, adapter_settings: ReflectionAdapterSettings) -> None:
        """Test that custom HNSW parameters (M, ef_construction) are used."""
        adapter = ReflectionDatabaseAdapterOneiric(settings=adapter_settings)
        await adapter.initialize()

        # Check index details to verify parameters
        # Note: Column is 'sql' not 'index_sql' (verified from duckdb_indexes() schema)
        result = adapter.conn.execute(
            """
            SELECT index_name, sql
            FROM duckdb_indexes()
            WHERE index_name LIKE '%hnsw%'
            """
        ).fetchall()

        assert len(result) >= 2

        # Verify SQL contains our custom parameters
        for row in result:
            # Column 1 is 'sql' (index creation statement)
            index_sql = row[1]
            # Check that our parameters are in the index definition
            if "M" in index_sql:
                assert "16" in index_sql or "M = 16" in index_sql
            if "ef_construction" in index_sql:
                assert "200" in index_sql or "ef_construction = 200" in index_sql

        await adapter.aclose()


class TestHNSWSearchPerformance:
    """Test that vector search works with HNSW indexes."""

    @pytest.fixture
    def temp_db_path(self, tmp_path: Path) -> Path:
        """Create temporary database path."""
        return tmp_path / "test_search.duckdb"

    @pytest.fixture
    async def populated_adapter(self, temp_db_path: Path) -> ReflectionDatabaseAdapterOneiric:
        """Create adapter with sample conversation data."""
        settings = ReflectionAdapterSettings(
            database_path=temp_db_path,
            collection_name="test_search",
            enable_hnsw_index=True,
        )

        adapter = ReflectionDatabaseAdapterOneiric(settings=settings)
        await adapter.initialize()

        # Add sample conversations with embeddings
        conversations = [
            "Python programming language",
            "Machine learning algorithms",
            "Web development with HTML",
            "Database management systems",
            "Cloud computing infrastructure",
        ]

        for conv in conversations:
            await adapter.store_conversation(conv, {"test": True})

        return adapter

    @pytest.mark.asyncio
    async def test_vector_search_with_hnsw(self, populated_adapter: ReflectionDatabaseAdapterOneiric) -> None:
        """Test that vector search works correctly with HNSW indexes."""
        # Search for similar conversations
        results = await populated_adapter.search_conversations(
            query="Python code", limit=5, threshold=0.0
        )

        # Should return results
        assert len(results) > 0

        # Check that results have expected structure
        for result in results:
            assert "id" in result
            assert "content" in result
            assert "score" in result
            assert isinstance(result["score"], float)
            assert 0.0 <= result["score"] <= 1.0

        # Top result should be about Python (highest similarity)
        assert "python" in results[0]["content"].lower()

        await populated_adapter.aclose()

    @pytest.mark.asyncio
    async def test_hnsw_ef_search_parameter_set(self, populated_adapter: ReflectionDatabaseAdapterOneiric) -> None:
        """Test that hnsw_ef_search parameter is set during search."""
        # Perform a search (should set hnsw_ef_search parameter)
        await populated_adapter.search_conversations("Python", limit=3)

        # Verify the parameter is set
        result = populated_adapter.conn.execute("SELECT * FROM duckdb_settings() WHERE name = 'hnsw_ef_search'").fetchone()

        # Parameter should be set to our configured value (64)
        assert result is not None
        assert result[1] == "64"  # value field

        await populated_adapter.aclose()


class TestHNSWGracefulFallback:
    """Test graceful fallback when VSS extension is unavailable."""

    @pytest.fixture
    def temp_db_path(self, tmp_path: Path) -> Path:
        """Create temporary database path."""
        return tmp_path / "test_fallback.duckdb"

    @pytest.mark.asyncio
    async def test_fallback_without_vss_extension(self, temp_db_path: Path) -> None:
        """Test that system falls back to array_cosine_similarity when VSS unavailable."""
        settings = ReflectionAdapterSettings(
            database_path=temp_db_path,
            collection_name="test_fallback",
            enable_hnsw_index=True,  # Try to enable HNSW
        )

        adapter = ReflectionDatabaseAdapterOneiric(settings=settings)

        # Initialize (should handle VSS extension gracefully)
        await adapter.initialize()

        # Add test data
        await adapter.store_conversation("Test conversation one", {"test": True})
        await adapter.store_conversation("Test conversation two", {"test": True})

        # Search should still work using array_cosine_similarity
        # Use threshold=0.0 to ensure we get results even with lower similarity
        results = await adapter.search_conversations("Test", limit=5, threshold=0.0)

        # Should return results using fallback method
        assert len(results) >= 2, f"Expected at least 2 results, got {len(results)}"

        # Results should have valid scores even without HNSW
        for result in results:
            assert "score" in result
            assert isinstance(result["score"], float)

        await adapter.aclose()

    @pytest.mark.asyncio
    async def test_hnsw_disabled_no_error(self, temp_db_path: Path) -> None:
        """Test that disabling HNSW doesn't cause any errors."""
        settings = ReflectionAdapterSettings(
            database_path=temp_db_path,
            collection_name="test_disabled",
            enable_hnsw_index=False,  # Explicitly disable
        )

        adapter = ReflectionDatabaseAdapterOneiric(settings=settings)
        await adapter.initialize()

        # Add and search data
        await adapter.store_conversation("Data without HNSW", {"test": True})
        # Use threshold=0.0 to ensure we get results
        results = await adapter.search_conversations("Data", limit=5, threshold=0.0)

        # Should work fine without HNSW
        assert len(results) >= 1, f"Expected at least 1 result, got {len(results)}"

        await adapter.aclose()


class TestHNSWConfigurationOptions:
    """Test various HNSW configuration options."""

    @pytest.fixture
    def temp_db_path(self, tmp_path: Path) -> Path:
        """Create temporary database path."""
        return tmp_path / "test_config.duckdb"

    @pytest.mark.asyncio
    async def test_custom_hnsw_parameters(self, temp_db_path: Path) -> None:
        """Test creating HNSW indexes with custom M and ef_construction values."""
        settings = ReflectionAdapterSettings(
            database_path=temp_db_path,
            collection_name="test_custom",
            enable_hnsw_index=True,
            hnsw_m=32,  # Custom M value
            hnsw_ef_construction=400,  # Custom ef_construction
        )

        adapter = ReflectionDatabaseAdapterOneiric(settings=settings)
        await adapter.initialize()

        # Verify indexes were created successfully
        # Note: The SQL returned by duckdb_indexes() may not include the WITH clause
        # The parameters are still used during index creation internally
        result = adapter.conn.execute(
            "SELECT index_name FROM duckdb_indexes() WHERE index_name LIKE '%_hnsw'"
        ).fetchall()

        # Should have created HNSW indexes
        assert len(result) >= 2, f"Expected at least 2 HNSW indexes, got {len(result)}"

        await adapter.aclose()

    @pytest.mark.asyncio
    async def test_hnsw_with_different_metrics(self, temp_db_path: Path) -> None:
        """Test HNSW indexes with different distance metrics."""
        # Test with L2 (Euclidean) distance
        # Note: DuckDB VSS uses "l2sq" for Euclidean distance, not "euclidean"
        settings_l2 = ReflectionAdapterSettings(
            database_path=temp_db_path,
            collection_name="test_l2",
            distance_metric="l2sq",  # DuckDB VSS uses "l2sq" for Euclidean distance
            enable_hnsw_index=True,
        )

        adapter = ReflectionDatabaseAdapterOneiric(settings=settings_l2)
        await adapter.initialize()

        # Verify index was created successfully
        result = adapter.conn.execute(
            "SELECT index_name FROM duckdb_indexes() WHERE index_name LIKE '%conv_embeddings_hnsw'"
        ).fetchone()

        assert result is not None
        # Column 0 is index_name
        assert "conv_embeddings_hnsw" in result[0]

        await adapter.aclose()


class TestHNSWPerformanceBasics:
    """Basic tests to verify HNSW functionality for performance improvements."""

    @pytest.fixture
    def temp_db_path(self, tmp_path: Path) -> Path:
        """Create temporary database path."""
        return tmp_path / "test_perf.duckdb"

    @pytest.mark.asyncio
    async def test_index_exists_check(self, temp_db_path: Path) -> None:
        """Test that we can check if HNSW indexes exist."""
        settings = ReflectionAdapterSettings(
            database_path=temp_db_path,
            collection_name="test_exist",
            enable_hnsw_index=True,
        )

        adapter = ReflectionDatabaseAdapterOneiric(settings=settings)
        await adapter.initialize()

        # Verify we can query for HNSW indexes
        result = adapter.conn.execute(
            """
            SELECT table_name, index_name
            FROM duckdb_indexes()
            WHERE index_name LIKE '%hnsw%'
            """
        ).fetchall()

        assert len(result) >= 2  # At least conversations and reflections

        await adapter.aclose()

    @pytest.mark.asyncio
    async def test_multiple_collections_independent_indexes(self, temp_db_path: Path) -> None:
        """Test that different collections get independent HNSW indexes."""
        # Create first collection
        settings1 = ReflectionAdapterSettings(
            database_path=temp_db_path,
            collection_name="collection1",
            enable_hnsw_index=True,
        )
        adapter1 = ReflectionDatabaseAdapterOneiric(settings=settings1)
        await adapter1.initialize()

        # Create second collection
        settings2 = ReflectionAdapterSettings(
            database_path=temp_db_path,
            collection_name="collection2",
            enable_hnsw_index=True,
        )
        adapter2 = ReflectionDatabaseAdapterOneiric(settings=settings2)
        await adapter2.initialize()

        # Check that both collections have their own HNSW indexes
        # Note: SELECT index_name means column 0 is the index name
        result = adapter1.conn.execute(
            "SELECT index_name FROM duckdb_indexes() WHERE index_name LIKE '%hnsw%'"
        ).fetchall()

        # Column 0 is index_name (from SELECT clause)
        collection1_indexes = [row[0] for row in result if "collection1" in row[0]]
        collection2_indexes = [row[0] for row in result if "collection2" in row[0]]

        # Each collection should have at least 2 HNSW indexes (conversations + reflections)
        assert len(collection1_indexes) >= 2
        assert len(collection2_indexes) >= 2

        await adapter1.aclose()
        await adapter2.aclose()


class TestHNSWWithRealData:
    """Integration tests with realistic data volumes."""

    @pytest.fixture
    def temp_db_path(self, tmp_path: Path) -> Path:
        """Create temporary database path."""
        return tmp_path / "test_realdata.duckdb"

    @pytest.mark.asyncio
    async def test_search_accuracy_with_hnsw(self, temp_db_path: Path) -> None:
        """Test that HNSW search maintains accuracy compared to baseline."""
        settings = ReflectionAdapterSettings(
            database_path=temp_db_path,
            collection_name="test_accuracy",
            enable_hnsw_index=True,
        )

        adapter = ReflectionDatabaseAdapterOneiric(settings=settings)
        await adapter.initialize()

        # Add related conversations
        await adapter.store_conversation("How to write Python functions", {"topic": "python"})
        await adapter.store_conversation("Python class inheritance tutorial", {"topic": "python"})
        await adapter.store_conversation("JavaScript function examples", {"topic": "javascript"})
        await adapter.store_conversation("Database query optimization", {"topic": "database"})

        # Search for Python-related content
        results = await adapter.search_conversations(
            "Python programming", limit=5, threshold=0.5
        )

        # Python results should rank higher than JavaScript
        python_results = [r for r in results if "python" in r["content"].lower()]
        js_results = [r for r in results if "javascript" in r["content"].lower()]

        # Python should appear before JavaScript in results
        if python_results and js_results:
            first_python_idx = next(i for i, r in enumerate(results) if "python" in r["content"].lower())
            first_js_idx = next(i for i, r in enumerate(results) if "javascript" in r["content"].lower())
            assert first_python_idx < first_js_idx, "Python results should rank higher"

        await adapter.aclose()