"""
Tests for cortex/hippo/tools/expand.py — hippo_expand tool.

Issue: #1118
"""

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 RetrievalEngine
from cortex.hippo.tools.expand import execute_expand


# ═══════════════════════════════════════════════════════════════════════════
# 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, retrieval):
    """Engine + retrieval with a 2-level summary DAG."""
    summary_store = engine.summary_store
    conv_id = engine._conversation.conversation_id

    # Ingest messages for leaf source
    for role, content in [
        ("user", "Let's deploy the app"),
        ("assistant", "Setting up fly.io deploy pipeline"),
        ("user", "What about the database migration?"),
        ("assistant", "Adding owner_id column migration"),
    ]:
        await engine.ingest_message(role=role, content=content)

    # Leaf summaries
    s1 = await summary_store.create_summary(
        conversation_id=conv_id,
        owner_id="test-owner",
        kind="leaf",
        depth=0,
        content="Deployed application to production via fly.io.",
        token_count=12,
    )
    s2 = await summary_store.create_summary(
        conversation_id=conv_id,
        owner_id="test-owner",
        kind="leaf",
        depth=0,
        content="Database schema migration adds owner_id column.",
        token_count=10,
    )

    # Condensed parent
    s3 = await summary_store.create_summary(
        conversation_id=conv_id,
        owner_id="test-owner",
        kind="condensed",
        depth=1,
        content="Sprint: deploy + db migration.",
        token_count=8,
        descendant_count=2,
        descendant_token_count=22,
    )

    # Link children
    await summary_store.add_summary_parent(s1.summary_id, s3.summary_id)
    await summary_store.add_summary_parent(s2.summary_id, s3.summary_id)

    return {
        "engine": engine,
        "retrieval": retrieval,
        "s1": s1,
        "s2": s2,
        "s3": s3,
        "conv_id": conv_id,
    }


# ═══════════════════════════════════════════════════════════════════════════
# Tests
# ═══════════════════════════════════════════════════════════════════════════


@pytest.mark.asyncio
async def test_expand_by_summary_ids(seeded):
    """Expand with explicit summaryIds returns children."""
    r = seeded["retrieval"]
    s3 = seeded["s3"]

    result = await execute_expand(
        params={"summaryIds": [s3.summary_id], "allConversations": True},
        retrieval=r,
        owner_id="test-owner",
    )

    text = result["content"][0]["text"]
    details = result["details"]

    assert details["expansionCount"] >= 1
    assert len(details["citedIds"]) >= 1
    # Should contain child summaries
    assert seeded["s1"].summary_id in text or seeded["s2"].summary_id in text


@pytest.mark.asyncio
async def test_expand_by_query(seeded):
    """Expand with query greps first, then expands matches."""
    r = seeded["retrieval"]

    result = await execute_expand(
        params={"query": "deploy", "allConversations": True},
        retrieval=r,
        owner_id="test-owner",
    )

    details = result["details"]
    # Should find at least one matching summary
    assert details["expansionCount"] >= 0  # May be 0 if FTS doesn't match
    assert "executionPath" in details


@pytest.mark.asyncio
async def test_max_depth_limit(seeded):
    """maxDepth=1 limits traversal to one level."""
    r = seeded["retrieval"]
    s3 = seeded["s3"]

    result = await execute_expand(
        params={"summaryIds": [s3.summary_id], "allConversations": True, "maxDepth": 1},
        retrieval=r,
        owner_id="test-owner",
    )

    details = result["details"]
    assert details["expansionCount"] >= 1


@pytest.mark.asyncio
async def test_token_cap_budget(seeded):
    """tokenCap limits total tokens in expansion."""
    r = seeded["retrieval"]
    s3 = seeded["s3"]

    # Very small cap — should trigger truncation
    result = await execute_expand(
        params={"summaryIds": [s3.summary_id], "allConversations": True, "tokenCap": 1},
        retrieval=r,
        owner_id="test-owner",
    )

    details = result["details"]
    assert details["truncated"] is True

    # Large cap — no truncation
    result2 = await execute_expand(
        params={"summaryIds": [s3.summary_id], "allConversations": True, "tokenCap": 100_000},
        retrieval=r,
        owner_id="test-owner",
    )

    details2 = result2["details"]
    assert details2["truncated"] is False


@pytest.mark.asyncio
async def test_include_messages(seeded):
    """includeMessages=true pulls source messages at leaf level."""
    r = seeded["retrieval"]
    s1 = seeded["s1"]

    # Link messages to leaf summary
    summary_store = seeded["engine"].summary_store
    messages = await seeded["engine"].conversation_store.get_messages(
        conversation_id=seeded["conv_id"],
        owner_id="test-owner",
        limit=4,
    )
    msg_ids = [msg.message_id for msg in messages[:2]]
    await summary_store.link_summary_messages(s1.summary_id, msg_ids)

    result = await execute_expand(
        params={
            "summaryIds": [s1.summary_id],
            "allConversations": True,
            "includeMessages": True,
        },
        retrieval=r,
        owner_id="test-owner",
    )

    text = result["content"][0]["text"]
    # Should have message content from source
    assert "deploy" in text.lower() or "fly.io" in text.lower() or "(no expansions)" not in text


@pytest.mark.asyncio
async def test_no_params_error(retrieval):
    """Missing both summaryIds and query returns error."""
    result = await execute_expand(
        params={"allConversations": True},
        retrieval=retrieval,
        owner_id="test-owner",
    )

    text = result["content"][0]["text"]
    assert "summaryIds or query must be provided" in text


@pytest.mark.asyncio
async def test_invalid_summary_ids_error(retrieval):
    """Invalid summary IDs (not sum_ prefix) are filtered out."""
    result = await execute_expand(
        params={"summaryIds": ["bad_id", "also_bad"], "allConversations": True},
        retrieval=retrieval,
        owner_id="test-owner",
    )

    text = result["content"][0]["text"]
    assert "No valid summary IDs" in text


@pytest.mark.asyncio
async def test_conversation_scope_mismatch(seeded):
    """Expand returns error when summaryIds belong to different conversation."""
    r = seeded["retrieval"]
    s1 = seeded["s1"]

    result = await execute_expand(
        params={"summaryIds": [s1.summary_id], "conversationId": 99999},
        retrieval=r,
        owner_id="test-owner",
    )

    text = result["content"][0]["text"]
    assert "outside conversation" in text
