"""
Tests for cortex/hippo/expansion.py — ExpansionOrchestrator.

Issue: #807
"""

from __future__ import annotations

import pytest
import pytest_asyncio

from cortex.hippo.config import HippoConfig
from cortex.hippo.engine import HippoEngine
from cortex.hippo.expansion import (
    ExpansionOrchestrator,
    ExpansionRequest,
    ExpansionResult,
    ExpansionEntry,
    ExpandQueryResult,
    execute_expand_query,
    resolve_expansion_token_cap,
)
from cortex.hippo.retrieval import ChildItem, MessageItem, RetrievalEngine


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


@pytest.fixture
def hippo_config() -> HippoConfig:
    return HippoConfig(
        enabled=True, db_path=":memory:", fresh_tail_count=4,
        max_tokens=10_000, max_expand_tokens=5_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 seeded_data(engine):
    """Engine with messages, leaf summaries, and a condensed summary."""
    # Create messages
    for role, content in [
        ("user", "Let's deploy the application to production"),
        ("assistant", "Setting up deploy pipeline via fly.io"),
        ("user", "What about database schema migration?"),
        ("assistant", "Adding owner_id column via migration v17"),
        ("user", "Fix the authentication bug too"),
        ("assistant", "Fixed auth bug - JWT validation now correct"),
    ]:
        await engine.ingest_message(role=role, content=content)

    summary_store = engine.summary_store
    conv_id = engine._conversation.conversation_id

    # Create leaf summaries
    leaf1 = await summary_store.create_summary(
        conversation_id=conv_id, owner_id="test-owner",
        kind="leaf", depth=0,
        content="Deployed app to production. DB migration adds owner_id.",
        token_count=15,
    )
    leaf2 = await summary_store.create_summary(
        conversation_id=conv_id, owner_id="test-owner",
        kind="leaf", depth=0,
        content="Fixed authentication bug. JWT token validation corrected.",
        token_count=12,
    )
    leaf3 = await summary_store.create_summary(
        conversation_id=conv_id, owner_id="test-owner",
        kind="leaf", depth=0,
        content="Discussed infrastructure scaling and monitoring setup.",
        token_count=10,
    )

    # Link messages
    await summary_store.link_summary_messages(leaf1.summary_id, [1, 2, 3, 4])
    await summary_store.link_summary_messages(leaf2.summary_id, [5, 6])

    # Create condensed summary
    condensed = await summary_store.create_summary(
        conversation_id=conv_id, owner_id="test-owner",
        kind="condensed", depth=1,
        content="Session: deploy, DB migration, auth fix, infra scaling.",
        token_count=12,
    )
    await summary_store.add_summary_parent(leaf1.summary_id, condensed.summary_id)
    await summary_store.add_summary_parent(leaf2.summary_id, condensed.summary_id)
    await summary_store.add_summary_parent(leaf3.summary_id, condensed.summary_id)

    return engine, leaf1, leaf2, leaf3, condensed


# ═══════════════════════════════════════════════════════════════════════════
# resolve_expansion_token_cap tests
# ═══════════════════════════════════════════════════════════════════════════


class TestResolveTokenCap:
    def test_none_uses_max(self):
        assert resolve_expansion_token_cap(None, 50000) == 50000

    def test_clamps_high(self):
        assert resolve_expansion_token_cap(100000, 50000) == 50000

    def test_clamps_low(self):
        assert resolve_expansion_token_cap(0, 50000) == 1

    def test_passes_through(self):
        assert resolve_expansion_token_cap(1000, 50000) == 1000


# ═══════════════════════════════════════════════════════════════════════════
# ExpansionOrchestrator.expand tests
# ═══════════════════════════════════════════════════════════════════════════


class TestExpand:
    @pytest.mark.asyncio
    async def test_expand_condensed_returns_children(self, seeded_data):
        engine, leaf1, leaf2, leaf3, condensed = seeded_data
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)
        orchestrator = ExpansionOrchestrator(retrieval, engine.config)

        result = await orchestrator.expand(ExpansionRequest(
            summary_ids=[condensed.summary_id],
            conversation_id=engine._conversation.conversation_id,
        ))
        assert len(result.expansions) == 1
        assert len(result.expansions[0].children) == 3
        assert condensed.summary_id in result.cited_ids

    @pytest.mark.asyncio
    async def test_expand_leaf_with_messages(self, seeded_data):
        engine, leaf1, leaf2, leaf3, condensed = seeded_data
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)
        orchestrator = ExpansionOrchestrator(retrieval, engine.config)

        result = await orchestrator.expand(ExpansionRequest(
            summary_ids=[leaf1.summary_id],
            conversation_id=engine._conversation.conversation_id,
            include_messages=True,
        ))
        assert len(result.expansions) == 1
        assert len(result.expansions[0].messages) == 4

    @pytest.mark.asyncio
    async def test_expand_multiple_ids(self, seeded_data):
        engine, leaf1, leaf2, leaf3, condensed = seeded_data
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)
        orchestrator = ExpansionOrchestrator(retrieval, engine.config)

        result = await orchestrator.expand(ExpansionRequest(
            summary_ids=[leaf1.summary_id, leaf2.summary_id],
            conversation_id=engine._conversation.conversation_id,
            include_messages=True,
        ))
        assert len(result.expansions) == 2
        assert leaf1.summary_id in result.cited_ids
        assert leaf2.summary_id in result.cited_ids

    @pytest.mark.asyncio
    async def test_expand_token_cap_truncates(self, seeded_data):
        engine, leaf1, leaf2, leaf3, condensed = seeded_data
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)
        orchestrator = ExpansionOrchestrator(retrieval, engine.config)

        result = await orchestrator.expand(ExpansionRequest(
            summary_ids=[condensed.summary_id],
            conversation_id=engine._conversation.conversation_id,
            token_cap=10,  # Very small — should truncate
        ))
        assert result.truncated is True
        # total_tokens should not wildly exceed the cap
        assert result.total_tokens <= 10 + 20  # small buffer for first child

    @pytest.mark.asyncio
    async def test_expand_global_token_cap_across_ids(self, seeded_data):
        """Token cap is GLOBAL across all summary IDs."""
        engine, leaf1, leaf2, leaf3, condensed = seeded_data
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)
        orchestrator = ExpansionOrchestrator(retrieval, engine.config)

        result = await orchestrator.expand(ExpansionRequest(
            summary_ids=[leaf1.summary_id, leaf2.summary_id, leaf3.summary_id],
            conversation_id=engine._conversation.conversation_id,
            include_messages=True,
            token_cap=20,  # Small cap shared across all IDs
        ))
        # Should not have expanded everything
        assert result.truncated is True

    @pytest.mark.asyncio
    async def test_expand_empty_ids(self, seeded_data):
        engine, *_ = seeded_data
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)
        orchestrator = ExpansionOrchestrator(retrieval, engine.config)

        result = await orchestrator.expand(ExpansionRequest(
            summary_ids=[],
            conversation_id=engine._conversation.conversation_id,
        ))
        assert len(result.expansions) == 0
        assert result.total_tokens == 0
        assert result.truncated is False


# ═══════════════════════════════════════════════════════════════════════════
# describe_and_expand tests
# ═══════════════════════════════════════════════════════════════════════════


class TestDescribeAndExpand:
    @pytest.mark.asyncio
    async def test_describe_and_expand_grep_first(self, seeded_data):
        engine, leaf1, leaf2, leaf3, condensed = seeded_data
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)
        orchestrator = ExpansionOrchestrator(retrieval, engine.config)

        result = await orchestrator.describe_and_expand(
            query="deploy",
            owner_id="test-owner",
            conversation_id=engine._conversation.conversation_id,
        )
        assert isinstance(result, ExpansionResult)
        # Should find at least one summary containing "deploy"
        assert len(result.expansions) >= 1

    @pytest.mark.asyncio
    async def test_describe_and_expand_no_results(self, seeded_data):
        engine, *_ = seeded_data
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)
        orchestrator = ExpansionOrchestrator(retrieval, engine.config)

        result = await orchestrator.describe_and_expand(
            query="nonexistent_xyz_term_99999",
            owner_id="test-owner",
        )
        assert len(result.expansions) == 0
        assert result.total_tokens == 0


# ═══════════════════════════════════════════════════════════════════════════
# distill_for_subagent tests
# ═══════════════════════════════════════════════════════════════════════════


class TestDistillForSubagent:
    def test_format_with_children(self):
        result = ExpansionResult(
            expansions=[
                ExpansionEntry(
                    summary_id="sum_abc123",
                    children=[
                        ChildItem(summary_id="sum_child1", kind="leaf",
                                  content="Deploy completed", token_count=50),
                        ChildItem(summary_id="sum_child2", kind="leaf",
                                  content="Auth fix applied", token_count=30),
                    ],
                ),
            ],
            cited_ids=["sum_abc123", "sum_child1", "sum_child2"],
            total_tokens=80,
            truncated=False,
        )
        output = ExpansionOrchestrator.distill_for_subagent(result)
        assert "## Expansion Results" in output
        assert "sum_abc123" in output
        assert "sum_child1" in output
        assert "Cited IDs for follow-up" in output
        assert "[Truncated: no]" in output

    def test_format_with_messages(self):
        result = ExpansionResult(
            expansions=[
                ExpansionEntry(
                    summary_id="sum_leaf1",
                    messages=[
                        MessageItem(message_id=1, role="user",
                                    content="Deploy please", token_count=3),
                    ],
                ),
            ],
            cited_ids=["sum_leaf1"],
            total_tokens=3,
            truncated=False,
        )
        output = ExpansionOrchestrator.distill_for_subagent(result)
        assert "msg#1" in output
        assert "user" in output

    def test_format_truncated(self):
        result = ExpansionResult(
            expansions=[], cited_ids=[], total_tokens=0, truncated=True,
        )
        output = ExpansionOrchestrator.distill_for_subagent(result)
        assert "[Truncated: yes]" in output

    def test_format_empty(self):
        result = ExpansionResult()
        output = ExpansionOrchestrator.distill_for_subagent(result)
        assert "## Expansion Results (0 summaries" in output
        assert "[Truncated: no]" in output


# ═══════════════════════════════════════════════════════════════════════════
# execute_expand_query tests
# ═══════════════════════════════════════════════════════════════════════════


class TestExpandQuery:
    @pytest.mark.asyncio
    async def test_expand_query_returns_cited_ids(self, seeded_data):
        engine, leaf1, leaf2, leaf3, condensed = seeded_data
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)
        orchestrator = ExpansionOrchestrator(retrieval, engine.config)

        async def mock_complete(**kwargs):
            return {"content": [{"type": "text", "text": "The deploy was to fly.io (sum_abc)."}]}

        result = await execute_expand_query(
            prompt="What was deployed?",
            orchestrator=orchestrator,
            complete_fn=mock_complete,
            config=engine.config,
            query="deploy",
            owner_id="test-owner",
            conversation_id=engine._conversation.conversation_id,
        )
        assert isinstance(result, ExpandQueryResult)
        assert len(result.cited_ids) >= 1
        assert "deploy" in result.answer.lower() or "fly" in result.answer.lower()

    @pytest.mark.asyncio
    async def test_expand_query_no_results(self, seeded_data):
        engine, *_ = seeded_data
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)
        orchestrator = ExpansionOrchestrator(retrieval, engine.config)

        async def mock_complete(**kwargs):
            return {"content": [{"type": "text", "text": "Nothing found."}]}

        result = await execute_expand_query(
            prompt="What about quantum computing?",
            orchestrator=orchestrator,
            complete_fn=mock_complete,
            config=engine.config,
            query="quantum_nonexistent_xyz",
            owner_id="test-owner",
        )
        assert "No summaries found" in result.answer

    @pytest.mark.asyncio
    async def test_expand_query_with_summary_ids(self, seeded_data):
        engine, leaf1, leaf2, leaf3, condensed = seeded_data
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)
        orchestrator = ExpansionOrchestrator(retrieval, engine.config)

        async def mock_complete(**kwargs):
            return {"content": [{"type": "text", "text": "Authentication was fixed."}]}

        result = await execute_expand_query(
            prompt="What was the auth fix?",
            orchestrator=orchestrator,
            complete_fn=mock_complete,
            config=engine.config,
            summary_ids=[leaf2.summary_id],
            owner_id="test-owner",
        )
        assert isinstance(result, ExpandQueryResult)
        assert leaf2.summary_id in result.cited_ids

    @pytest.mark.asyncio
    async def test_expand_query_llm_failure(self, seeded_data):
        engine, leaf1, leaf2, leaf3, condensed = seeded_data
        retrieval = RetrievalEngine(engine.conversation_store, engine.summary_store)
        orchestrator = ExpansionOrchestrator(retrieval, engine.config)

        async def failing_complete(**kwargs):
            raise RuntimeError("LLM unavailable")

        result = await execute_expand_query(
            prompt="What happened?",
            orchestrator=orchestrator,
            complete_fn=failing_complete,
            config=engine.config,
            query="deploy",
            owner_id="test-owner",
            conversation_id=engine._conversation.conversation_id,
        )
        assert "failed" in result.answer.lower() or "found" in result.answer.lower()
        assert len(result.cited_ids) >= 1
