"""
Tests for cortex/hippo/retrieval.py — RetrievalEngine search and inspect.

Issue: #806
"""

from __future__ import annotations

import pytest
import pytest_asyncio

from cortex.hippo.config import HippoConfig
from cortex.hippo.engine import HippoEngine
from cortex.hippo.retrieval import (
    DescribeResult,
    ExpandInput,
    GrepInput,
    GrepResult,
    RetrievalEngine,
    cosine_similarity,
    pack_floats,
    unpack_floats,
)
from cortex.hippo.store.fts_utils import sanitize_fts5_query


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


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


@pytest_asyncio.fixture
async def engine(hippo_config):
    eng = HippoEngine(config=hippo_config)
    await eng.initialize(db_path=":memory:", owner_id="test-owner", conversation_id="test-session")
    yield eng
    await eng.shutdown()


@pytest_asyncio.fixture
async def retrieval(engine):
    return RetrievalEngine(
        conversation_store=engine.conversation_store,
        summary_store=engine.summary_store,
    )


@pytest_asyncio.fixture
async def seeded_engine(engine):
    """Engine with some messages and summaries pre-loaded."""
    # Insert messages
    for i, (role, content) in enumerate([
        ("user", "Let's deploy the application to production"),
        ("assistant", "I'll set up the deploy pipeline using fly.io"),
        ("user", "What about the database schema migration?"),
        ("assistant", "The database schema needs a migration to add the owner_id column"),
        ("user", "Let's also fix the authentication bug"),
        ("assistant", "Fixed the auth bug by validating JWT tokens properly"),
    ]):
        await engine.ingest_message(role=role, content=content)

    # Create summaries
    summary_store = engine.summary_store
    s1 = await summary_store.create_summary(
        conversation_id=engine._conversation.conversation_id,
        owner_id="test-owner",
        kind="leaf",
        depth=0,
        content="Deployed application to production via fly.io. Database schema migration adds owner_id column.",
        token_count=20,
    )
    s2 = await summary_store.create_summary(
        conversation_id=engine._conversation.conversation_id,
        owner_id="test-owner",
        kind="leaf",
        depth=0,
        content="Fixed authentication bug by validating JWT tokens. Auth flow now works correctly.",
        token_count=18,
    )

    # Link messages to summaries
    await summary_store.link_summary_messages(s1.summary_id, [1, 2, 3, 4])
    await summary_store.link_summary_messages(s2.summary_id, [5, 6])

    return engine, s1, s2


# ═══════════════════════════════════════════════════════════════════════════
# FTS5 utility tests
# ═══════════════════════════════════════════════════════════════════════════


class TestFTS5Utils:
    def test_sanitize_normal_query(self):
        assert sanitize_fts5_query("deploy pipeline") == "deploy pipeline"

    def test_sanitize_special_chars(self):
        result = sanitize_fts5_query('this (won\'t "work":easily)')
        assert result is not None
        assert "(" not in result
        assert '"' not in result
        assert ":" not in result

    def test_sanitize_empty_query(self):
        assert sanitize_fts5_query("") is None

    def test_sanitize_single_char_tokens(self):
        assert sanitize_fts5_query("a b c") is None

    def test_sanitize_mixed_tokens(self):
        result = sanitize_fts5_query("a deploy b to c production")
        assert result == "deploy to production"


# ═══════════════════════════════════════════════════════════════════════════
# Vector utility tests
# ═══════════════════════════════════════════════════════════════════════════


class TestVectorUtils:
    def test_pack_unpack_roundtrip(self):
        original = [1.0, 2.5, -3.14, 0.0]
        packed = pack_floats(original)
        unpacked = unpack_floats(packed)
        assert len(unpacked) == len(original)
        for a, b in zip(original, unpacked):
            assert abs(a - b) < 1e-5

    def test_cosine_similarity_identical(self):
        v = [1.0, 2.0, 3.0]
        assert abs(cosine_similarity(v, v) - 1.0) < 1e-6

    def test_cosine_similarity_orthogonal(self):
        a = [1.0, 0.0]
        b = [0.0, 1.0]
        assert abs(cosine_similarity(a, b)) < 1e-6

    def test_cosine_similarity_opposite(self):
        a = [1.0, 0.0]
        b = [-1.0, 0.0]
        assert abs(cosine_similarity(a, b) - (-1.0)) < 1e-6

    def test_cosine_similarity_zero_vector(self):
        assert cosine_similarity([0.0, 0.0], [1.0, 2.0]) == 0.0

    def test_cosine_similarity_empty(self):
        assert cosine_similarity([], []) == 0.0

    def test_cosine_similarity_length_mismatch(self):
        assert cosine_similarity([1.0], [1.0, 2.0]) == 0.0


# ═══════════════════════════════════════════════════════════════════════════
# Grep tests
# ═══════════════════════════════════════════════════════════════════════════


class TestGrep:
    @pytest.mark.asyncio
    async def test_grep_regex_messages(self, seeded_engine):
        engine, s1, s2 = seeded_engine
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)

        result = await retrieval.grep(GrepInput(
            pattern="deploy",
            mode="regex",
            scope="messages",
            owner_id="test-owner",
        ))
        assert isinstance(result, GrepResult)
        assert result.total_matches >= 1
        assert all("deploy" in m.content.lower() for m in result.messages)

    @pytest.mark.asyncio
    async def test_grep_regex_summaries(self, seeded_engine):
        engine, s1, s2 = seeded_engine
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)

        result = await retrieval.grep(GrepInput(
            pattern="JWT",
            mode="regex",
            scope="summaries",
            owner_id="test-owner",
        ))
        assert result.total_matches >= 1
        assert any("JWT" in s.content for s in result.summaries)

    @pytest.mark.asyncio
    async def test_grep_both_scopes(self, seeded_engine):
        engine, s1, s2 = seeded_engine
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)

        result = await retrieval.grep(GrepInput(
            pattern="deploy",
            mode="regex",
            scope="both",
            owner_id="test-owner",
        ))
        # Should have results from both messages and summaries
        assert result.total_matches >= 2

    @pytest.mark.asyncio
    async def test_grep_fts5_messages(self, seeded_engine):
        engine, s1, s2 = seeded_engine
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)

        result = await retrieval.grep(GrepInput(
            pattern="database schema",
            mode="full_text",
            scope="messages",
            owner_id="test-owner",
        ))
        assert isinstance(result, GrepResult)
        assert result.total_matches >= 1

    @pytest.mark.asyncio
    async def test_grep_fts5_summaries(self, seeded_engine):
        engine, s1, s2 = seeded_engine
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)

        result = await retrieval.grep(GrepInput(
            pattern="authentication",
            mode="full_text",
            scope="summaries",
            owner_id="test-owner",
        ))
        assert result.total_matches >= 1

    @pytest.mark.asyncio
    async def test_grep_no_results(self, seeded_engine):
        engine, s1, s2 = seeded_engine
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)

        result = await retrieval.grep(GrepInput(
            pattern="nonexistent_xyz_term_12345",
            mode="regex",
            scope="both",
            owner_id="test-owner",
        ))
        assert result.total_matches == 0

    @pytest.mark.asyncio
    async def test_grep_semantic_falls_back_without_embeddings(self, seeded_engine):
        engine, s1, s2 = seeded_engine
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store, embedding_provider=None)

        # Should fall back to full_text
        result = await retrieval.grep(GrepInput(
            pattern="deploy",
            mode="semantic",
            scope="both",
            owner_id="test-owner",
        ))
        assert isinstance(result, GrepResult)

    @pytest.mark.asyncio
    async def test_grep_with_conversation_filter(self, seeded_engine):
        engine, s1, s2 = seeded_engine
        conv_id = engine._conversation.conversation_id
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)

        result = await retrieval.grep(GrepInput(
            pattern="deploy",
            mode="regex",
            scope="messages",
            conversation_id=conv_id,
            owner_id="test-owner",
        ))
        assert all(m.conversation_id == conv_id for m in result.messages)

    @pytest.mark.asyncio
    async def test_grep_limit(self, seeded_engine):
        engine, s1, s2 = seeded_engine
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)

        result = await retrieval.grep(GrepInput(
            pattern="the",
            mode="regex",
            scope="messages",
            owner_id="test-owner",
            limit=2,
        ))
        assert len(result.messages) <= 2


# ═══════════════════════════════════════════════════════════════════════════
# Describe tests
# ═══════════════════════════════════════════════════════════════════════════


class TestDescribe:
    @pytest.mark.asyncio
    async def test_describe_summary_returns_lineage(self, seeded_engine):
        engine, s1, s2 = seeded_engine
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)

        result = await retrieval.describe(s1.summary_id, owner_id="test-owner")
        assert result is not None
        assert result.type == "summary"
        assert result.summary is not None
        assert result.summary["summary_id"] == s1.summary_id
        assert isinstance(result.summary["parent_ids"], list)
        assert isinstance(result.summary["child_ids"], list)
        assert isinstance(result.summary["message_ids"], list)
        assert isinstance(result.summary["subtree"], list)

    @pytest.mark.asyncio
    async def test_describe_summary_has_messages(self, seeded_engine):
        engine, s1, s2 = seeded_engine
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)

        result = await retrieval.describe(s1.summary_id, owner_id="test-owner")
        assert len(result.summary["message_ids"]) == 4  # linked to msgs 1-4

    @pytest.mark.asyncio
    async def test_describe_nonexistent(self, seeded_engine):
        engine, s1, s2 = seeded_engine
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)

        result = await retrieval.describe("sum_nonexistent1234", owner_id="test-owner")
        assert result is None

    @pytest.mark.asyncio
    async def test_describe_unknown_prefix(self, seeded_engine):
        engine, s1, s2 = seeded_engine
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)

        result = await retrieval.describe("unknown_abc123", owner_id="test-owner")
        assert result is None

    @pytest.mark.asyncio
    async def test_describe_file_nonexistent(self, seeded_engine):
        engine, s1, s2 = seeded_engine
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)

        result = await retrieval.describe("file_doesntexist", owner_id="test-owner")
        assert result is None


# ═══════════════════════════════════════════════════════════════════════════
# Semantic search tests
# ═══════════════════════════════════════════════════════════════════════════


class MockEmbeddingProvider:
    """Mock embedding provider that returns deterministic vectors."""

    async def embed(self, text: str) -> list[float]:
        # Simple hash-based vector for testing
        h = hash(text) & 0xFFFFFFFF
        # Create a small 4-dim vector based on text features
        has_deploy = 1.0 if "deploy" in text.lower() else 0.0
        has_auth = 1.0 if "auth" in text.lower() else 0.0
        has_db = 1.0 if "database" in text.lower() or "schema" in text.lower() else 0.0
        base = (h % 100) / 100.0
        return [has_deploy + base, has_auth + base, has_db + base, base]


class TestSemanticSearch:
    @pytest.mark.asyncio
    async def test_semantic_search_returns_results(self, seeded_engine):
        engine, s1, s2 = seeded_engine
        provider = MockEmbeddingProvider()
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store, embedding_provider=provider)

        # Store embeddings on summaries
        s1_vec = await provider.embed(s1.content)
        s2_vec = await provider.embed(s2.content)
        await engine._db.conn.execute(
            "UPDATE hippo_summaries SET embedding = ? WHERE summary_id = ?",
            (pack_floats(s1_vec), s1.summary_id),
        )
        await engine._db.conn.execute(
            "UPDATE hippo_summaries SET embedding = ? WHERE summary_id = ?",
            (pack_floats(s2_vec), s2.summary_id),
        )
        await engine._db.conn.commit()

        results = await retrieval.semantic_search(
            query="deploy to production",
            owner_id="test-owner",
            score_threshold=0.0,  # Low threshold for test
        )
        assert len(results) >= 1
        assert all(r.score >= 0.0 for r in results)

    @pytest.mark.asyncio
    async def test_semantic_search_no_provider(self, seeded_engine):
        engine, s1, s2 = seeded_engine
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store, embedding_provider=None)

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

    @pytest.mark.asyncio
    async def test_grep_semantic_with_provider(self, seeded_engine):
        engine, s1, s2 = seeded_engine
        provider = MockEmbeddingProvider()
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store, embedding_provider=provider)

        # Store embeddings
        s1_vec = await provider.embed(s1.content)
        await engine._db.conn.execute(
            "UPDATE hippo_summaries SET embedding = ? WHERE summary_id = ?",
            (pack_floats(s1_vec), s1.summary_id),
        )
        await engine._db.conn.commit()

        result = await retrieval.grep(GrepInput(
            pattern="deploy to production",
            mode="semantic",
            scope="summaries",
            owner_id="test-owner",
        ))
        assert isinstance(result, GrepResult)


# ═══════════════════════════════════════════════════════════════════════════
# Expand tests (single summary)
# ═══════════════════════════════════════════════════════════════════════════


class TestExpand:
    @pytest.mark.asyncio
    async def test_expand_leaf_with_messages(self, seeded_engine):
        engine, s1, s2 = seeded_engine
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)

        result = await retrieval.expand(ExpandInput(
            summary_id=s1.summary_id,
            depth=1,
            include_messages=True,
        ))
        assert len(result.messages) == 4  # s1 linked to msgs 1-4
        assert result.estimated_tokens > 0

    @pytest.mark.asyncio
    async def test_expand_leaf_without_messages(self, seeded_engine):
        engine, s1, s2 = seeded_engine
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)

        result = await retrieval.expand(ExpandInput(
            summary_id=s1.summary_id,
            depth=1,
            include_messages=False,
        ))
        # Leaf with no include_messages returns no children (it's a leaf)
        assert len(result.children) == 0
        assert len(result.messages) == 0

    @pytest.mark.asyncio
    async def test_expand_condensed_returns_children(self, seeded_engine):
        engine, s1, s2 = seeded_engine
        summary_store = engine.summary_store

        # Create a condensed summary that parents s1 and s2
        condensed = await summary_store.create_summary(
            conversation_id=engine._conversation.conversation_id,
            owner_id="test-owner",
            kind="condensed",
            depth=1,
            content="Full session: deploy, DB migration, auth fix.",
            token_count=12,
        )
        # Link: s1 and s2 are children of condensed
        await summary_store.add_summary_parent(s1.summary_id, condensed.summary_id)
        await summary_store.add_summary_parent(s2.summary_id, condensed.summary_id)

        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)
        result = await retrieval.expand(ExpandInput(
            summary_id=condensed.summary_id,
            depth=2,
            include_messages=True,
        ))
        assert len(result.children) == 2  # s1 and s2

    @pytest.mark.asyncio
    async def test_expand_token_cap(self, seeded_engine):
        engine, s1, s2 = seeded_engine
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)

        result = await retrieval.expand(ExpandInput(
            summary_id=s1.summary_id,
            depth=1,
            include_messages=True,
            token_cap=5,  # Very small cap
        ))
        # Should truncate before getting all messages
        assert result.truncated is True
        assert result.estimated_tokens <= 5 + 50  # Small buffer for first message

    @pytest.mark.asyncio
    async def test_expand_depth_zero(self, seeded_engine):
        engine, s1, s2 = seeded_engine
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)

        result = await retrieval.expand(ExpandInput(
            summary_id=s1.summary_id,
            depth=0,
            include_messages=True,
        ))
        # depth=0 means no expansion
        assert len(result.children) == 0
        assert len(result.messages) == 0


# ═══════════════════════════════════════════════════════════════════════════
# Subtree CTE tests
# ═══════════════════════════════════════════════════════════════════════════


class TestSubtreeCTE:
    @pytest.mark.asyncio
    async def test_subtree_linear_chain(self, seeded_engine):
        """Test subtree with linear chain: root -> A -> B."""
        engine, s1, s2 = seeded_engine
        summary_store = engine.summary_store
        conv_id = engine._conversation.conversation_id

        root = await summary_store.create_summary(
            conversation_id=conv_id, owner_id="test-owner",
            kind="condensed", depth=2, content="root", token_count=5,
        )
        mid = await summary_store.create_summary(
            conversation_id=conv_id, owner_id="test-owner",
            kind="condensed", depth=1, content="mid", token_count=5,
        )
        leaf = await summary_store.create_summary(
            conversation_id=conv_id, owner_id="test-owner",
            kind="leaf", depth=0, content="leaf", token_count=5,
        )

        await summary_store.add_summary_parent(mid.summary_id, root.summary_id)
        await summary_store.add_summary_parent(leaf.summary_id, mid.summary_id)

        subtree = await summary_store.get_summary_subtree(root.summary_id)
        ids = [n["summary_id"] for n in subtree]
        assert root.summary_id in ids
        assert mid.summary_id in ids
        assert leaf.summary_id in ids
        assert len(subtree) == 3

        # Verify depths
        depth_map = {n["summary_id"]: n["depth_from_root"] for n in subtree}
        assert depth_map[root.summary_id] == 0
        assert depth_map[mid.summary_id] == 1
        assert depth_map[leaf.summary_id] == 2

    @pytest.mark.asyncio
    async def test_subtree_branching_tree(self, seeded_engine):
        """Test subtree with branching: root -> (A, B)."""
        engine, s1, s2 = seeded_engine
        summary_store = engine.summary_store
        conv_id = engine._conversation.conversation_id

        root = await summary_store.create_summary(
            conversation_id=conv_id, owner_id="test-owner",
            kind="condensed", depth=1, content="root", token_count=5,
        )
        child_a = await summary_store.create_summary(
            conversation_id=conv_id, owner_id="test-owner",
            kind="leaf", depth=0, content="child_a", token_count=5,
        )
        child_b = await summary_store.create_summary(
            conversation_id=conv_id, owner_id="test-owner",
            kind="leaf", depth=0, content="child_b", token_count=5,
        )

        await summary_store.add_summary_parent(child_a.summary_id, root.summary_id)
        await summary_store.add_summary_parent(child_b.summary_id, root.summary_id)

        subtree = await summary_store.get_summary_subtree(root.summary_id)
        assert len(subtree) == 3

    @pytest.mark.asyncio
    async def test_subtree_diamond_dag_no_duplicates(self, seeded_engine):
        """Test diamond DAG: root -> (A, B) -> leaf. Leaf should appear once (UNION dedup)."""
        engine, s1, s2 = seeded_engine
        summary_store = engine.summary_store
        conv_id = engine._conversation.conversation_id

        root = await summary_store.create_summary(
            conversation_id=conv_id, owner_id="test-owner",
            kind="condensed", depth=2, content="root", token_count=5,
        )
        mid_a = await summary_store.create_summary(
            conversation_id=conv_id, owner_id="test-owner",
            kind="condensed", depth=1, content="mid_a", token_count=5,
        )
        mid_b = await summary_store.create_summary(
            conversation_id=conv_id, owner_id="test-owner",
            kind="condensed", depth=1, content="mid_b", token_count=5,
        )
        leaf = await summary_store.create_summary(
            conversation_id=conv_id, owner_id="test-owner",
            kind="leaf", depth=0, content="shared_leaf", token_count=5,
        )

        # root -> mid_a, root -> mid_b
        await summary_store.add_summary_parent(mid_a.summary_id, root.summary_id)
        await summary_store.add_summary_parent(mid_b.summary_id, root.summary_id)
        # mid_a -> leaf, mid_b -> leaf (diamond)
        await summary_store.add_summary_parent(leaf.summary_id, mid_a.summary_id)
        await summary_store.add_summary_parent(leaf.summary_id, mid_b.summary_id)

        subtree = await summary_store.get_summary_subtree(root.summary_id)
        ids = [n["summary_id"] for n in subtree]

        # Should have 4 unique nodes, not 5 (leaf appears once due to UNION)
        assert len(ids) == 4
        assert ids.count(leaf.summary_id) == 1

    @pytest.mark.asyncio
    async def test_subtree_depth_cap(self, seeded_engine):
        """Test that depth cap is enforced."""
        engine, s1, s2 = seeded_engine
        summary_store = engine.summary_store
        conv_id = engine._conversation.conversation_id

        # Build a chain of depth 5
        prev_id = None
        all_ids = []
        for i in range(5):
            s = await summary_store.create_summary(
                conversation_id=conv_id, owner_id="test-owner",
                kind="condensed" if i < 4 else "leaf",
                depth=4 - i, content=f"node_{i}", token_count=5,
            )
            all_ids.append(s.summary_id)
            if prev_id is not None:
                await summary_store.add_summary_parent(s.summary_id, prev_id)
            prev_id = s.summary_id

        # With max_depth=2, should only get 3 levels (root + 2)
        subtree = await summary_store.get_summary_subtree(all_ids[0], max_depth=2)
        assert len(subtree) == 3  # root, depth 1, depth 2

    @pytest.mark.asyncio
    async def test_subtree_single_node(self, seeded_engine):
        """Test subtree of a leaf with no children."""
        engine, s1, s2 = seeded_engine
        summary_store = engine.summary_store

        subtree = await summary_store.get_summary_subtree(s1.summary_id)
        assert len(subtree) == 1
        assert subtree[0]["summary_id"] == s1.summary_id
        assert subtree[0]["depth_from_root"] == 0
