"""
Tests for batch query methods in summary_store and conversation_store,
and BFS expansion in retrieval.

Issue: #1187
"""

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 ExpandInput, RetrievalEngine


OWNER = "test-owner"


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


@pytest_asyncio.fixture
async def eng(batch_config: HippoConfig) -> HippoEngine:
    e = HippoEngine(config=batch_config)
    await e.initialize(db_path=":memory:", owner_id=OWNER, conversation_id="test-session")
    yield e
    await e.shutdown()


async def _create_summary(eng, kind="leaf", depth=0, content="test summary", token_count=50):
    return await eng.summary_store.create_summary(
        conversation_id=eng._conversation.conversation_id,
        owner_id=OWNER,
        kind=kind,
        depth=depth,
        content=content,
        token_count=token_count,
    )


@pytest.mark.asyncio
async def test_get_records_batch(eng: HippoEngine):
    """Batch summary fetch returns correct records."""
    s1 = await _create_summary(eng, content="summary one")
    s2 = await _create_summary(eng, content="summary two")
    s3 = await _create_summary(eng, content="summary three")

    results = await eng.summary_store.get_records_batch(
        [s1.summary_id, s2.summary_id, s3.summary_id]
    )
    ids = {r.summary_id for r in results}
    assert ids == {s1.summary_id, s2.summary_id, s3.summary_id}


@pytest.mark.asyncio
async def test_get_records_batch_empty(eng: HippoEngine):
    """Empty input returns empty list."""
    results = await eng.summary_store.get_records_batch([])
    assert results == []


@pytest.mark.asyncio
async def test_get_children_batch(eng: HippoEngine):
    """Batch fetch returns correct children per parent."""
    parent1 = await _create_summary(eng, kind="condensed", depth=1, content="parent 1")
    parent2 = await _create_summary(eng, kind="condensed", depth=1, content="parent 2")
    child1 = await _create_summary(eng, content="child of parent 1")
    child2 = await _create_summary(eng, content="child of parent 2")
    child3 = await _create_summary(eng, content="another child of parent 1")

    # Note: in hippo_summary_parents, (summary_id, parent_summary_id)
    # means summary_id is the child and parent_summary_id is the parent
    await eng.summary_store.add_summary_parent(child1.summary_id, parent1.summary_id)
    await eng.summary_store.add_summary_parent(child3.summary_id, parent1.summary_id)
    await eng.summary_store.add_summary_parent(child2.summary_id, parent2.summary_id)

    result = await eng.summary_store.get_children_records_batch(
        [parent1.summary_id, parent2.summary_id]
    )

    assert len(result[parent1.summary_id]) == 2
    assert len(result[parent2.summary_id]) == 1
    p1_child_ids = {c.summary_id for c in result[parent1.summary_id]}
    assert child1.summary_id in p1_child_ids
    assert child3.summary_id in p1_child_ids
    assert result[parent2.summary_id][0].summary_id == child2.summary_id


@pytest.mark.asyncio
async def test_get_children_batch_empty(eng: HippoEngine):
    """Empty input returns empty dict."""
    result = await eng.summary_store.get_children_records_batch([])
    assert result == {}


@pytest.mark.asyncio
async def test_get_source_messages_batch(eng: HippoEngine):
    """Batch source message IDs fetch works."""
    conv_id = eng._conversation.conversation_id
    msg1 = await eng.conversation_store.create_message(conv_id, OWNER, "user", "msg 1")
    msg2 = await eng.conversation_store.create_message(conv_id, OWNER, "assistant", "msg 2")
    msg3 = await eng.conversation_store.create_message(conv_id, OWNER, "user", "msg 3")

    s1 = await _create_summary(eng, content="leaf 1")
    s2 = await _create_summary(eng, content="leaf 2")

    await eng.summary_store.link_summary_messages(s1.summary_id, [msg1.message_id, msg2.message_id])
    await eng.summary_store.link_summary_messages(s2.summary_id, [msg3.message_id])

    result = await eng.summary_store.get_source_message_ids_batch(
        [s1.summary_id, s2.summary_id]
    )

    assert set(result[s1.summary_id]) == {msg1.message_id, msg2.message_id}
    assert result[s2.summary_id] == [msg3.message_id]


@pytest.mark.asyncio
async def test_get_messages_by_ids_batch(eng: HippoEngine):
    """Batch message fetch via get_messages_by_ids works."""
    conv_id = eng._conversation.conversation_id
    msg1 = await eng.conversation_store.create_message(conv_id, OWNER, "user", "hello")
    msg2 = await eng.conversation_store.create_message(conv_id, OWNER, "assistant", "world")

    result = await eng.conversation_store.get_messages_by_ids(
        [msg1.message_id, msg2.message_id]
    )
    assert msg1.message_id in result
    assert msg2.message_id in result
    assert result[msg1.message_id].content == "hello"
    assert result[msg2.message_id].content == "world"


@pytest.mark.asyncio
async def test_expand_bfs_matches_recursive(eng: HippoEngine):
    """BFS expansion should produce same children as recursive for a simple DAG."""
    conv_id = eng._conversation.conversation_id

    # Build: root -> [child1, child2], child1 -> [leaf1], child2 -> [leaf2]
    root = await _create_summary(eng, kind="condensed", depth=2, content="root", token_count=30)
    child1 = await _create_summary(eng, kind="condensed", depth=1, content="child1", token_count=20)
    child2 = await _create_summary(eng, kind="condensed", depth=1, content="child2", token_count=20)
    leaf1 = await _create_summary(eng, kind="leaf", depth=0, content="leaf1", token_count=10)
    leaf2 = await _create_summary(eng, kind="leaf", depth=0, content="leaf2", token_count=10)

    await eng.summary_store.add_summary_parent(child1.summary_id, root.summary_id)
    await eng.summary_store.add_summary_parent(child2.summary_id, root.summary_id)
    await eng.summary_store.add_summary_parent(leaf1.summary_id, child1.summary_id)
    await eng.summary_store.add_summary_parent(leaf2.summary_id, child2.summary_id)

    retrieval = RetrievalEngine(
        conversation_store=eng.conversation_store,
        summary_store=eng.summary_store,
    )

    result = await retrieval.expand(ExpandInput(
        summary_id=root.summary_id,
        depth=3,
        include_messages=False,
    ))

    # Should find child1, child2, leaf1, leaf2
    child_ids = {c.summary_id for c in result.children}
    assert child1.summary_id in child_ids
    assert child2.summary_id in child_ids
    assert leaf1.summary_id in child_ids
    assert leaf2.summary_id in child_ids


@pytest.mark.asyncio
async def test_expand_bfs_respects_token_cap(eng: HippoEngine):
    """BFS expansion should respect token budget."""
    root = await _create_summary(eng, kind="condensed", depth=1, content="root", token_count=10)
    child1 = await _create_summary(eng, kind="leaf", depth=0, content="child1", token_count=100)
    child2 = await _create_summary(eng, kind="leaf", depth=0, content="child2", token_count=100)

    await eng.summary_store.add_summary_parent(child1.summary_id, root.summary_id)
    await eng.summary_store.add_summary_parent(child2.summary_id, root.summary_id)

    retrieval = RetrievalEngine(
        conversation_store=eng.conversation_store,
        summary_store=eng.summary_store,
    )

    result = await retrieval.expand(ExpandInput(
        summary_id=root.summary_id,
        depth=2,
        include_messages=False,
        token_cap=150,
    ))

    # Only one child should fit (100 tokens each, cap is 150)
    assert len(result.children) == 1
    assert result.truncated is True


@pytest.mark.asyncio
async def test_expand_bfs_with_messages(eng: HippoEngine):
    """BFS expansion should fetch source messages for leaf nodes."""
    conv_id = eng._conversation.conversation_id
    msg = await eng.conversation_store.create_message(conv_id, OWNER, "user", "original message")

    root = await _create_summary(eng, kind="condensed", depth=1, content="root", token_count=10)
    leaf = await _create_summary(eng, kind="leaf", depth=0, content="leaf summary", token_count=10)

    await eng.summary_store.add_summary_parent(leaf.summary_id, root.summary_id)
    await eng.summary_store.link_summary_messages(leaf.summary_id, [msg.message_id])

    retrieval = RetrievalEngine(
        conversation_store=eng.conversation_store,
        summary_store=eng.summary_store,
    )

    result = await retrieval.expand(ExpandInput(
        summary_id=root.summary_id,
        depth=2,
        include_messages=True,
    ))

    # Should have the leaf as a child and the message
    assert len(result.children) == 1
    assert result.children[0].summary_id == leaf.summary_id
    assert len(result.messages) == 1
    assert result.messages[0].content == "original message"
