"""
Tests for sqlite-vec ANN vector search integration in Hippo.

Issue: #1184
"""

from __future__ import annotations

import math
import struct
from unittest.mock import AsyncMock, patch

import pytest
import pytest_asyncio

from cortex.hippo.config import HippoConfig
from cortex.hippo.db.connection import HippoDatabase
from cortex.hippo.retrieval import (
    RetrievalEngine,
    SemanticSearchResult,
    cosine_similarity,
    pack_floats,
    unpack_floats,
)
from cortex.hippo.store.summary_store import SummaryStore
from cortex.hippo.store.conversation_store import ConversationStore


# ═══════════════════════════════════════════════════════════════════════════
# Fixtures
# ═══════════════════════════════════════════════════════════════════════════

DIM = 4  # small dimension for tests


def make_vec(values: list[float]) -> bytes:
    """Pack floats into a vec blob."""
    return pack_floats(values)


def make_unit_vec(idx: int, dim: int = DIM) -> list[float]:
    """Create a unit vector with 1.0 at position idx, 0.0 elsewhere."""
    v = [0.0] * dim
    v[idx % dim] = 1.0
    return v


def make_similar_vec(base: list[float], noise: float = 0.1) -> list[float]:
    """Create a vector similar to base by adding small noise."""
    return [x + noise for x in base]


@pytest.fixture
def config() -> HippoConfig:
    return HippoConfig(
        enabled=True,
        db_path=":memory:",
        embedding_dim=DIM,
        vec_enabled=True,
        fresh_tail_count=4,
        max_tokens=10_000,
    )


@pytest.fixture
def config_no_vec() -> HippoConfig:
    return HippoConfig(
        enabled=True,
        db_path=":memory:",
        embedding_dim=DIM,
        vec_enabled=False,
        fresh_tail_count=4,
        max_tokens=10_000,
    )


@pytest_asyncio.fixture
async def db(config) -> HippoDatabase:
    database = HippoDatabase(config)
    await database.initialize()
    yield database
    await database.close()


@pytest_asyncio.fixture
async def db_no_vec(config_no_vec) -> HippoDatabase:
    database = HippoDatabase(config_no_vec)
    await database.initialize()
    yield database
    await database.close()


@pytest_asyncio.fixture
async def stores(db):
    """Return (conversation_store, summary_store) with vec enabled."""
    conv_store = ConversationStore(db.conn)
    sum_store = SummaryStore(db.conn, vec_available=db.vec_available)
    return conv_store, sum_store


@pytest_asyncio.fixture
async def stores_no_vec(db_no_vec):
    """Return (conversation_store, summary_store) without vec."""
    conv_store = ConversationStore(db_no_vec.conn)
    sum_store = SummaryStore(db_no_vec.conn, vec_available=False)
    return conv_store, sum_store


async def _create_conversation(conv_store, owner_id="test-owner"):
    """Helper: create a conversation and return its id."""
    conv = await conv_store.get_or_create_conversation(owner_id, "test-session")
    return conv.conversation_id


async def _create_summary_with_embedding(
    sum_store, conn, conv_id, owner_id, content, embedding_vec, kind="leaf",
):
    """Helper: create a summary and set its embedding + vec sync."""
    summary = await sum_store.create_summary(
        conversation_id=conv_id,
        owner_id=owner_id,
        kind=kind,
        depth=0,
        content=content,
        token_count=len(content) // 4,
    )
    embedding_blob = make_vec(embedding_vec)
    await sum_store.update_summary_embedding(
        summary.summary_id, owner_id, embedding_blob,
    )
    return summary


# ═══════════════════════════════════════════════════════════════════════════
# Extension & Table Tests
# ═══════════════════════════════════════════════════════════════════════════


class TestVecExtensionLoading:
    """Test sqlite-vec extension loads correctly."""

    @pytest.mark.asyncio
    async def test_vec_extension_loads(self, db):
        """sqlite-vec should be available and loaded."""
        assert db.vec_available is True
        # Verify vec_version() works
        async with db.conn.execute("SELECT vec_version()") as cur:
            row = await cur.fetchone()
            assert row is not None
            assert row[0]  # version string is non-empty

    @pytest.mark.asyncio
    async def test_vec_disabled_by_config(self, db_no_vec):
        """When vec_enabled=False, extension should not be loaded."""
        assert db_no_vec.vec_available is False

    @pytest.mark.asyncio
    async def test_vec_table_created(self, db):
        """vec0 virtual table should exist after initialization."""
        async with db.conn.execute(
            "SELECT name FROM sqlite_master WHERE name = 'hippo_vec_summaries'"
        ) as cur:
            row = await cur.fetchone()
            assert row is not None
            assert row[0] == "hippo_vec_summaries"


# ═══════════════════════════════════════════════════════════════════════════
# Migration Tests
# ═══════════════════════════════════════════════════════════════════════════


class TestVecMigration:
    """Test the add_vec_index migration."""

    @pytest.mark.asyncio
    async def test_vec_migration_populates_existing(self, config):
        """Migration should populate vec table from existing embeddings."""
        from cortex.hippo.db.migrations.add_vec_index import migrate, needs_migration

        # Create a fresh DB with vec loaded but without the vec table
        database = HippoDatabase(config)
        await database.initialize()

        # Insert embeddings directly
        conv_store = ConversationStore(database.conn)
        conv = await conv_store.get_or_create_conversation("owner1", "sess1")
        sum_store = SummaryStore(database.conn, vec_available=False)  # skip vec sync
        s = await sum_store.create_summary(
            conversation_id=conv.conversation_id,
            owner_id="owner1", kind="leaf", depth=0,
            content="test content", token_count=5,
        )
        emb = make_vec([1.0, 0.0, 0.0, 0.0])
        await database.conn.execute(
            "UPDATE hippo_summaries SET embedding = ? WHERE summary_id = ?",
            (emb, s.summary_id),
        )
        await database.conn.commit()

        # Drop the vec table to simulate pre-migration state
        await database.conn.executescript("DROP TABLE IF EXISTS hippo_vec_summaries")

        # Recreate it via migration
        assert await needs_migration(database.conn) is True
        await migrate(database.conn, embedding_dim=DIM)
        assert await needs_migration(database.conn) is False

        # Verify the embedding was populated
        async with database.conn.execute(
            "SELECT summary_id FROM hippo_vec_summaries"
        ) as cur:
            rows = await cur.fetchall()
            assert len(rows) == 1
            assert rows[0][0] == s.summary_id

        await database.close()


# ═══════════════════════════════════════════════════════════════════════════
# Vec Search Tests
# ═══════════════════════════════════════════════════════════════════════════


class TestVecSearch:
    """Test ANN vector search via sqlite-vec."""

    @pytest.mark.asyncio
    async def test_vec_search_returns_results(self, db, stores):
        """Basic vec search should return matching results."""
        conv_store, sum_store = stores
        conv_id = await _create_conversation(conv_store)

        vec_a = [1.0, 0.0, 0.0, 0.0]
        vec_b = [0.0, 1.0, 0.0, 0.0]
        await _create_summary_with_embedding(
            sum_store, db.conn, conv_id, "test-owner", "Alpha summary", vec_a,
        )
        await _create_summary_with_embedding(
            sum_store, db.conn, conv_id, "test-owner", "Beta summary", vec_b,
        )

        provider = AsyncMock()
        provider.embed = AsyncMock(return_value=vec_a)

        retrieval = RetrievalEngine(
            conversation_store=conv_store,
            summary_store=sum_store,
            embedding_provider=provider,
            db=db,
        )
        results = await retrieval.semantic_search(
            query="alpha", owner_id="test-owner", limit=5,
        )
        assert len(results) >= 1
        assert results[0].score > 0.9  # near-identical to vec_a

    @pytest.mark.asyncio
    async def test_vec_search_filters_by_owner(self, db, stores):
        """Vec search should only return results for the specified owner."""
        conv_store, sum_store = stores

        # Two owners sharing same conversation
        conv_id = await _create_conversation(conv_store, "alice")

        vec = [1.0, 0.0, 0.0, 0.0]
        await _create_summary_with_embedding(
            sum_store, db.conn, conv_id, "alice", "Alice's note", vec,
        )
        # Create Bob's conversation and summary
        conv_bob = await conv_store.get_or_create_conversation("bob", "bob-sess")
        await _create_summary_with_embedding(
            sum_store, db.conn, conv_bob.conversation_id, "bob", "Bob's note", vec,
        )

        provider = AsyncMock()
        provider.embed = AsyncMock(return_value=vec)

        retrieval = RetrievalEngine(
            conversation_store=conv_store,
            summary_store=sum_store,
            embedding_provider=provider,
            db=db,
        )

        alice_results = await retrieval.semantic_search(
            query="test", owner_id="alice", limit=10,
        )
        bob_results = await retrieval.semantic_search(
            query="test", owner_id="bob", limit=10,
        )

        alice_ids = {r.summary_id for r in alice_results}
        bob_ids = {r.summary_id for r in bob_results}
        assert alice_ids.isdisjoint(bob_ids)
        assert len(alice_results) == 1
        assert len(bob_results) == 1

    @pytest.mark.asyncio
    async def test_vec_search_filters_by_conversation(self, db, stores):
        """Vec search with conversation_id should filter results."""
        conv_store, sum_store = stores
        conv_id_1 = (await conv_store.get_or_create_conversation("owner", "s1")).conversation_id
        conv_id_2 = (await conv_store.get_or_create_conversation("owner", "s2")).conversation_id

        vec = [1.0, 0.0, 0.0, 0.0]
        s1 = await _create_summary_with_embedding(
            sum_store, db.conn, conv_id_1, "owner", "Conv 1 content", vec,
        )
        s2 = await _create_summary_with_embedding(
            sum_store, db.conn, conv_id_2, "owner", "Conv 2 content", vec,
        )

        provider = AsyncMock()
        provider.embed = AsyncMock(return_value=vec)

        retrieval = RetrievalEngine(
            conversation_store=conv_store,
            summary_store=sum_store,
            embedding_provider=provider,
            db=db,
        )

        results = await retrieval.semantic_search(
            query="test", owner_id="owner", conversation_id=conv_id_1, limit=10,
        )
        result_ids = {r.summary_id for r in results}
        assert s1.summary_id in result_ids
        assert s2.summary_id not in result_ids

    @pytest.mark.asyncio
    async def test_vec_search_score_threshold(self, db, stores):
        """Results below score_threshold should be filtered out."""
        conv_store, sum_store = stores
        conv_id = await _create_conversation(conv_store)

        # Orthogonal vectors → cosine similarity = 0
        vec_a = [1.0, 0.0, 0.0, 0.0]
        vec_b = [0.0, 1.0, 0.0, 0.0]

        await _create_summary_with_embedding(
            sum_store, db.conn, conv_id, "test-owner", "Orthogonal", vec_b,
        )

        provider = AsyncMock()
        provider.embed = AsyncMock(return_value=vec_a)

        retrieval = RetrievalEngine(
            conversation_store=conv_store,
            summary_store=sum_store,
            embedding_provider=provider,
            db=db,
        )

        results = await retrieval.semantic_search(
            query="test", owner_id="test-owner", limit=5, score_threshold=0.5,
        )
        # Orthogonal vectors → similarity ~0 → filtered out
        assert len(results) == 0

    @pytest.mark.asyncio
    async def test_vec_search_empty_table(self, db, stores):
        """Vec search on empty table should return empty list, not crash."""
        conv_store, sum_store = stores

        provider = AsyncMock()
        provider.embed = AsyncMock(return_value=[1.0, 0.0, 0.0, 0.0])

        retrieval = RetrievalEngine(
            conversation_store=conv_store,
            summary_store=sum_store,
            embedding_provider=provider,
            db=db,
        )

        results = await retrieval.semantic_search(
            query="test", owner_id="test-owner", limit=5,
        )
        assert results == []


# ═══════════════════════════════════════════════════════════════════════════
# Parity & Fallback Tests
# ═══════════════════════════════════════════════════════════════════════════


class TestVecParity:
    """Test parity between vec and brute-force search."""

    @pytest.mark.asyncio
    async def test_vec_search_parity_with_brute_force(self, db, stores):
        """Vec and brute-force should return same results (within float precision)."""
        conv_store, sum_store = stores
        conv_id = await _create_conversation(conv_store)

        vecs = [
            [0.9, 0.1, 0.0, 0.0],
            [0.1, 0.9, 0.0, 0.0],
            [0.5, 0.5, 0.0, 0.0],
        ]
        for i, v in enumerate(vecs):
            await _create_summary_with_embedding(
                sum_store, db.conn, conv_id, "test-owner", f"Summary {i}", v,
            )

        query_vec = [0.8, 0.2, 0.0, 0.0]
        provider = AsyncMock()
        provider.embed = AsyncMock(return_value=query_vec)

        retrieval = RetrievalEngine(
            conversation_store=conv_store,
            summary_store=sum_store,
            embedding_provider=provider,
            db=db,
        )

        # Vec search
        vec_results = await retrieval._vec_semantic_search(
            query_vec, "test-owner", None, 10, 0.0,
        )
        # Brute-force search
        bf_results = await retrieval._brute_force_semantic_search(
            query_vec, "test-owner", None, 10, 0.0,
        )

        # Same IDs in same order
        vec_ids = [r.summary_id for r in vec_results]
        bf_ids = [r.summary_id for r in bf_results]
        assert vec_ids == bf_ids

        # Scores within float tolerance
        for vr, br in zip(vec_results, bf_results):
            assert abs(vr.score - br.score) < 0.01, (
                f"Score mismatch for {vr.summary_id}: vec={vr.score}, bf={br.score}"
            )

    @pytest.mark.asyncio
    async def test_vec_fallback_when_unavailable(self, db_no_vec, stores_no_vec):
        """When vec is not available, should fall back to brute-force."""
        conv_store, sum_store = stores_no_vec
        conv_id = await _create_conversation(conv_store)

        # Store embedding directly (no vec sync since vec is disabled)
        summary = await sum_store.create_summary(
            conversation_id=conv_id,
            owner_id="test-owner",
            kind="leaf", depth=0,
            content="Fallback test", token_count=5,
        )
        emb = make_vec([1.0, 0.0, 0.0, 0.0])
        await db_no_vec.conn.execute(
            "UPDATE hippo_summaries SET embedding = ? WHERE summary_id = ?",
            (emb, summary.summary_id),
        )
        await db_no_vec.conn.commit()

        provider = AsyncMock()
        provider.embed = AsyncMock(return_value=[1.0, 0.0, 0.0, 0.0])

        retrieval = RetrievalEngine(
            conversation_store=conv_store,
            summary_store=sum_store,
            embedding_provider=provider,
            db=db_no_vec,
        )

        results = await retrieval.semantic_search(
            query="test", owner_id="test-owner", limit=5,
        )
        assert len(results) == 1
        assert results[0].summary_id == summary.summary_id


# ═══════════════════════════════════════════════════════════════════════════
# Sync Tests
# ═══════════════════════════════════════════════════════════════════════════


class TestVecSync:
    """Test that vec index stays in sync with summary store operations."""

    @pytest.mark.asyncio
    async def test_vec_sync_on_create(self, db, stores):
        """New summary with embedding should sync to vec table."""
        conv_store, sum_store = stores
        conv_id = await _create_conversation(conv_store)

        vec = [1.0, 0.0, 0.0, 0.0]
        summary = await _create_summary_with_embedding(
            sum_store, db.conn, conv_id, "test-owner", "Synced summary", vec,
        )

        # Verify it's in the vec table
        async with db.conn.execute(
            "SELECT summary_id FROM hippo_vec_summaries WHERE summary_id = ?",
            (summary.summary_id,),
        ) as cur:
            row = await cur.fetchone()
            assert row is not None
            assert row[0] == summary.summary_id

    @pytest.mark.asyncio
    async def test_vec_sync_on_delete(self, db, stores):
        """Deleted summary should be removed from vec table."""
        conv_store, sum_store = stores
        conv_id = await _create_conversation(conv_store)

        vec = [1.0, 0.0, 0.0, 0.0]
        summary = await _create_summary_with_embedding(
            sum_store, db.conn, conv_id, "test-owner", "To be deleted", vec,
        )

        # Verify it exists
        async with db.conn.execute(
            "SELECT COUNT(*) FROM hippo_vec_summaries WHERE summary_id = ?",
            (summary.summary_id,),
        ) as cur:
            assert (await cur.fetchone())[0] == 1

        # Delete it
        await sum_store.delete_summary(summary.summary_id)

        # Verify it's gone
        async with db.conn.execute(
            "SELECT COUNT(*) FROM hippo_vec_summaries WHERE summary_id = ?",
            (summary.summary_id,),
        ) as cur:
            assert (await cur.fetchone())[0] == 0
