"""
Test: touch_compaction_heartbeat is called during compaction loops.

Verifies B-006 fix: heartbeat_fn is wired into CompactionEngine and
called during compact_leaf, compact_full_sweep, and compact_until_under.
"""

from __future__ import annotations

import asyncio
import pytest
from unittest.mock import AsyncMock, MagicMock, patch

from cortex.hippo.compaction.engine import CompactionEngine
from cortex.hippo.compaction.types import CompactionResult, PassResult
from cortex.hippo.config import HippoConfig


@pytest.fixture
def config():
    return HippoConfig(
        enabled=True,
        max_rounds=3,
        context_threshold=0.8,
        incremental_condensed_max_depth=0,  # disable incremental condensed for leaf test
    )


@pytest.fixture
def heartbeat_fn():
    return AsyncMock()


@pytest.fixture
def stores():
    conv_store = AsyncMock()
    sum_store = AsyncMock()
    # get_context_token_count returns decreasing values to show progress
    sum_store.get_context_token_count = AsyncMock(side_effect=[
        1000, 800, 600, 400, 200, 100, 50, 50, 50, 50, 50,
    ])
    return conv_store, sum_store


class TestHeartbeatWiredIntoCompactLeaf:
    """Heartbeat is called after the leaf pass in compact_leaf."""

    @pytest.mark.asyncio
    async def test_heartbeat_called_during_compact_leaf(self, config, heartbeat_fn):
        conv_store = AsyncMock()
        sum_store = AsyncMock()
        # tokens_before, then tokens_after_leaf
        sum_store.get_context_token_count = AsyncMock(side_effect=[1000, 800])

        engine = CompactionEngine(
            conversation_store=conv_store,
            summary_store=sum_store,
            config=config,
            heartbeat_fn=heartbeat_fn,
        )

        fake_chunk = MagicMock()
        fake_chunk.items = [MagicMock(), MagicMock()]

        fake_pass = PassResult(
            summary_id="sum_test_1",
            level="normal",
            content="test summary",
            tokens_before=500,
            tokens_after=300,
        )

        with patch(
            "cortex.hippo.compaction.engine._select_oldest_leaf_chunk",
            new_callable=AsyncMock,
            return_value=fake_chunk,
        ), patch(
            "cortex.hippo.compaction.engine._resolve_prior_leaf_summary_context",
            new_callable=AsyncMock,
            return_value=None,
        ), patch(
            "cortex.hippo.compaction.engine._leaf_pass",
            new_callable=AsyncMock,
            return_value=fake_pass,
        ), patch(
            "cortex.hippo.compaction.engine._persist_compaction_events",
            new_callable=AsyncMock,
        ), patch(
            "cortex.hippo.compaction.engine._resolve_owner_id",
            new_callable=AsyncMock,
            return_value="owner1",
        ):
            result = await engine.compact_leaf(
                conversation_id=42,
                token_budget=2000,
                summarize=AsyncMock(),
                force=True,
            )

        assert result.action_taken
        # Heartbeat must have been called at least once with conversation_id=42
        heartbeat_fn.assert_called_with(42)
        assert heartbeat_fn.call_count >= 1


class TestHeartbeatWiredIntoCompactFullSweep:
    """Heartbeat is called after each leaf and condensed pass in compact_full_sweep."""

    @pytest.mark.asyncio
    async def test_heartbeat_called_during_full_sweep(self, config, heartbeat_fn):
        conv_store = AsyncMock()
        sum_store = AsyncMock()
        # Token counts: before, after leaf1, after leaf2 (no items), after_leaf final,
        # then condensed phase returns no candidate
        call_count = 0

        async def token_side_effect(conv_id):
            nonlocal call_count
            call_count += 1
            # Decreasing tokens to show progress, then stabilize
            values = [1000, 800, 600, 600, 600, 600, 600]
            return values[min(call_count - 1, len(values) - 1)]

        sum_store.get_context_token_count = AsyncMock(side_effect=token_side_effect)

        engine = CompactionEngine(
            conversation_store=conv_store,
            summary_store=sum_store,
            config=config,
            heartbeat_fn=heartbeat_fn,
        )

        fake_chunk = MagicMock()
        fake_chunk.items = [MagicMock(), MagicMock()]
        empty_chunk = MagicMock()
        empty_chunk.items = []

        fake_pass = PassResult(
            summary_id="sum_test_1",
            level="normal",
            content="test summary",
            tokens_before=500,
            tokens_after=300,
        )

        leaf_call_count = 0

        async def select_leaf_side_effect(*args, **kwargs):
            nonlocal leaf_call_count
            leaf_call_count += 1
            if leaf_call_count <= 2:
                return fake_chunk
            return empty_chunk

        with patch(
            "cortex.hippo.compaction.engine._select_oldest_leaf_chunk",
            new_callable=AsyncMock,
            side_effect=select_leaf_side_effect,
        ), patch(
            "cortex.hippo.compaction.engine._resolve_prior_leaf_summary_context",
            new_callable=AsyncMock,
            return_value=None,
        ), patch(
            "cortex.hippo.compaction.engine._leaf_pass",
            new_callable=AsyncMock,
            return_value=fake_pass,
        ), patch(
            "cortex.hippo.compaction.engine._select_shallowest_condensation_candidate",
            new_callable=AsyncMock,
            return_value=None,  # No condensation candidates
        ), patch(
            "cortex.hippo.compaction.engine._persist_compaction_events",
            new_callable=AsyncMock,
        ), patch(
            "cortex.hippo.compaction.engine._resolve_owner_id",
            new_callable=AsyncMock,
            return_value="owner1",
        ):
            result = await engine.compact_full_sweep(
                conversation_id=42,
                token_budget=2000,
                summarize=AsyncMock(),
                force=True,
            )

        assert result.action_taken
        # 2 leaf passes → 2 heartbeat calls (one after each leaf pass)
        assert heartbeat_fn.call_count >= 2
        heartbeat_fn.assert_called_with(42)


class TestHeartbeatWiredIntoCompactUntilUnder:
    """Heartbeat is called in compact_until_under loop."""

    @pytest.mark.asyncio
    async def test_heartbeat_called_during_compact_until_under(self, heartbeat_fn):
        config = HippoConfig(
            enabled=True,
            max_rounds=5,
            context_threshold=0.8,
            incremental_condensed_max_depth=0,
        )
        conv_store = AsyncMock()
        sum_store = AsyncMock()

        # Token counts decrease over sweep rounds
        token_values = [1000, 900, 800, 700, 600, 500, 400, 300, 200, 100]
        token_idx = 0

        async def token_side_effect(conv_id):
            nonlocal token_idx
            val = token_values[min(token_idx, len(token_values) - 1)]
            token_idx += 1
            return val

        sum_store.get_context_token_count = AsyncMock(side_effect=token_side_effect)

        engine = CompactionEngine(
            conversation_store=conv_store,
            summary_store=sum_store,
            config=config,
            heartbeat_fn=heartbeat_fn,
        )

        # Mock compact_full_sweep to return a result that shows action taken
        sweep_results = [
            CompactionResult(action_taken=True, tokens_before=1000, tokens_after=500, created_summary_ids=["s1"], rounds=1),
            CompactionResult(action_taken=False, tokens_before=500, tokens_after=500, created_summary_ids=[], rounds=0),
        ]
        sweep_idx = 0

        async def mock_sweep(*args, **kwargs):
            nonlocal sweep_idx
            r = sweep_results[min(sweep_idx, len(sweep_results) - 1)]
            sweep_idx += 1
            return r

        engine.compact_full_sweep = AsyncMock(side_effect=mock_sweep)

        result = await engine.compact_until_under(
            conversation_id=42,
            token_budget=2000,
            target_tokens=400,
            current_tokens=1000,
            summarize=AsyncMock(),
        )

        # Heartbeat should be called at least once (after each sweep round)
        assert heartbeat_fn.call_count >= 1
        heartbeat_fn.assert_called_with(42)


class TestHeartbeatNotCalledWhenNone:
    """When no heartbeat_fn is provided, no error occurs."""

    @pytest.mark.asyncio
    async def test_no_heartbeat_fn_no_error(self):
        config = HippoConfig(
            enabled=True,
            max_rounds=3,
            incremental_condensed_max_depth=0,
        )
        conv_store = AsyncMock()
        sum_store = AsyncMock()
        sum_store.get_context_token_count = AsyncMock(side_effect=[1000, 800])

        # No heartbeat_fn passed
        engine = CompactionEngine(
            conversation_store=conv_store,
            summary_store=sum_store,
            config=config,
        )

        fake_chunk = MagicMock()
        fake_chunk.items = [MagicMock()]
        fake_pass = PassResult(summary_id="sum_1", level="normal", content="test", tokens_before=500, tokens_after=300)

        with patch(
            "cortex.hippo.compaction.engine._select_oldest_leaf_chunk",
            new_callable=AsyncMock,
            return_value=fake_chunk,
        ), patch(
            "cortex.hippo.compaction.engine._resolve_prior_leaf_summary_context",
            new_callable=AsyncMock,
            return_value=None,
        ), patch(
            "cortex.hippo.compaction.engine._leaf_pass",
            new_callable=AsyncMock,
            return_value=fake_pass,
        ), patch(
            "cortex.hippo.compaction.engine._persist_compaction_events",
            new_callable=AsyncMock,
        ), patch(
            "cortex.hippo.compaction.engine._resolve_owner_id",
            new_callable=AsyncMock,
            return_value="owner1",
        ):
            # Should not raise
            result = await engine.compact_leaf(
                conversation_id=42,
                token_budget=2000,
                summarize=AsyncMock(),
                force=True,
            )
            assert result.action_taken


class TestHeartbeatCallbackErrorHandled:
    """If heartbeat_fn raises, compaction continues."""

    @pytest.mark.asyncio
    async def test_heartbeat_error_does_not_break_compaction(self):
        config = HippoConfig(
            enabled=True,
            max_rounds=3,
            incremental_condensed_max_depth=0,
        )
        conv_store = AsyncMock()
        sum_store = AsyncMock()
        sum_store.get_context_token_count = AsyncMock(side_effect=[1000, 800])

        failing_heartbeat = AsyncMock(side_effect=Exception("DB error"))

        engine = CompactionEngine(
            conversation_store=conv_store,
            summary_store=sum_store,
            config=config,
            heartbeat_fn=failing_heartbeat,
        )

        fake_chunk = MagicMock()
        fake_chunk.items = [MagicMock()]
        fake_pass = PassResult(summary_id="sum_1", level="normal", content="test", tokens_before=500, tokens_after=300)

        with patch(
            "cortex.hippo.compaction.engine._select_oldest_leaf_chunk",
            new_callable=AsyncMock,
            return_value=fake_chunk,
        ), patch(
            "cortex.hippo.compaction.engine._resolve_prior_leaf_summary_context",
            new_callable=AsyncMock,
            return_value=None,
        ), patch(
            "cortex.hippo.compaction.engine._leaf_pass",
            new_callable=AsyncMock,
            return_value=fake_pass,
        ), patch(
            "cortex.hippo.compaction.engine._persist_compaction_events",
            new_callable=AsyncMock,
        ), patch(
            "cortex.hippo.compaction.engine._resolve_owner_id",
            new_callable=AsyncMock,
            return_value="owner1",
        ):
            # Should NOT raise despite heartbeat error
            result = await engine.compact_leaf(
                conversation_id=42,
                token_budget=2000,
                summarize=AsyncMock(),
                force=True,
            )
            assert result.action_taken
            failing_heartbeat.assert_called_with(42)
