"""Tests for automatic dream cycle in EpisodeDigestor."""

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

import pytest

from neo_cortex.digestor import EpisodeDigestor
from neo_cortex.models import DreamResult


class _Stop(BaseException):
    """Escapes the try/except Exception inside digestor.run()."""


@pytest.fixture
def cortex_with_dream():
    cortex = AsyncMock()
    cortex.dream.return_value = DreamResult(
        status="completed", dream_count=1, cluster=["m1"], surprisal=["m2"]
    )
    return cortex


@pytest.fixture(autouse=True)
def _no_sleep():
    """Patch asyncio.sleep so digestor.run() doesn't block."""
    with patch("asyncio.sleep", new_callable=AsyncMock):
        yield


@pytest.mark.asyncio
async def test_needs_dream_starts_false():
    """_needs_dream flag starts False."""
    cortex = AsyncMock()
    log = MagicMock()
    digestor = EpisodeDigestor(log=log, cortex=cortex)
    assert digestor._needs_dream is False


@pytest.mark.asyncio
async def test_needs_dream_set_on_work(cortex_with_dream):
    """_needs_dream becomes True when _process_all returns True."""
    log = MagicMock()
    digestor = EpisodeDigestor(log=log, cortex=cortex_with_dream)

    call_count = 0

    async def one_work_then_stop():
        nonlocal call_count
        call_count += 1
        if call_count > 1:
            raise _Stop("stop")
        return True  # did work

    digestor._process_all = one_work_then_stop

    with pytest.raises(_Stop):
        await digestor.run()

    assert digestor._needs_dream is True


@pytest.mark.asyncio
async def test_dream_not_called_without_prior_work(cortex_with_dream):
    """Dream does NOT fire on idle if no work was done."""
    log = MagicMock()
    digestor = EpisodeDigestor(log=log, cortex=cortex_with_dream)

    call_count = 0

    async def idle_then_stop():
        nonlocal call_count
        call_count += 1
        if call_count > 1:
            raise _Stop("stop")
        return False  # no work

    digestor._process_all = idle_then_stop

    with pytest.raises(_Stop):
        await digestor.run()

    cortex_with_dream.dream.assert_not_called()


@pytest.mark.asyncio
async def test_dream_called_after_work_then_idle(cortex_with_dream):
    """Dream fires when transitioning work→idle."""
    log = MagicMock()
    digestor = EpisodeDigestor(log=log, cortex=cortex_with_dream)

    call_count = 0

    async def work_then_idle_then_stop():
        nonlocal call_count
        call_count += 1
        if call_count == 1:
            return True   # work
        if call_count == 2:
            return False  # idle → dream fires here
        raise _Stop("stop")

    digestor._process_all = work_then_idle_then_stop

    with pytest.raises(_Stop):
        await digestor.run()

    cortex_with_dream.dream.assert_called_once()
    assert digestor._needs_dream is False  # reset after dream


@pytest.mark.asyncio
async def test_dream_exception_does_not_crash(cortex_with_dream):
    """dream() exception doesn't crash the digestor loop."""
    cortex_with_dream.dream.side_effect = Exception("dream exploded")
    log = MagicMock()
    digestor = EpisodeDigestor(log=log, cortex=cortex_with_dream)
    digestor._needs_dream = True

    call_count = 0

    async def idle_then_stop():
        nonlocal call_count
        call_count += 1
        if call_count > 1:
            raise _Stop("stop")
        return False

    digestor._process_all = idle_then_stop

    with pytest.raises(_Stop):
        await digestor.run()

    cortex_with_dream.dream.assert_called_once()
    assert digestor._needs_dream is False  # reset even on failure


@pytest.mark.asyncio
async def test_dream_not_called_twice_on_consecutive_idles(cortex_with_dream):
    """Dream fires once after work, not again on next idle without new work."""
    log = MagicMock()
    digestor = EpisodeDigestor(log=log, cortex=cortex_with_dream)

    call_count = 0

    async def work_idle_idle_stop():
        nonlocal call_count
        call_count += 1
        if call_count == 1:
            return True   # work
        if call_count <= 3:
            return False  # two idle cycles
        raise _Stop("stop")

    digestor._process_all = work_idle_idle_stop

    with pytest.raises(_Stop):
        await digestor.run()

    # Dream should be called exactly once (after first idle, not second)
    assert cortex_with_dream.dream.call_count == 1
