"""Tests for EpisodeDigestor — token-budgeted batch processing."""

import time
from unittest.mock import AsyncMock

import pytest

import neo_cortex.config as cortex_config
from neo_cortex.conversation_log import ConversationLog
from neo_cortex.digestor import EpisodeDigestor, CHARS_PER_TOKEN
from neo_cortex.distiller import KnowledgeEntry
from neo_cortex.models import ConversationAppendRequest, ConversationEntry


def make_entry(
    id: int = 0,
    role: str = "user",
    content: str = "hello",
    session_id: str = "sess-001",
    timestamp: float | None = None,
) -> ConversationEntry:
    return ConversationEntry(
        id=id,
        timestamp=timestamp or time.time(),
        session_id=session_id,
        role=role,
        content=content,
    )


def _make_log(tmp_path, entries_data: list[dict]) -> ConversationLog:
    """Create a log with given entries."""
    db = str(tmp_path / "log.db")
    log = ConversationLog(db)
    for e in entries_data:
        log.append(ConversationAppendRequest(
            session_id=e.get("session_id", "sess-001"),
            role=e["role"],
            content=e["content"],
            timestamp=e.get("timestamp", time.time() - 120),
        ))
    return log


def _make_multi_topic_log(tmp_path) -> ConversationLog:
    """Create a log with 3 distinct topics, all old timestamps."""
    db = str(tmp_path / "log.db")
    log = ConversationLog(db)
    old = time.time() - 120

    # Topic 1: CORS fix
    log.append(ConversationAppendRequest(session_id="s1", role="user", content="Fix CORS on neo-reloaded", timestamp=old))
    log.append(ConversationAppendRequest(session_id="s1", role="assistant", content="Added app.UseCors() to pipeline", timestamp=old + 1))
    log.append(ConversationAppendRequest(session_id="s1", role="user", content="test it", timestamp=old + 2))
    log.append(ConversationAppendRequest(session_id="s1", role="assistant", content="132 tests pass", timestamp=old + 3))

    # Topic 2: systemd deploy
    log.append(ConversationAppendRequest(session_id="s1", role="user", content="Deploy to systemd", timestamp=old + 4))
    log.append(ConversationAppendRequest(session_id="s1", role="assistant", content="Created service file, enabled, started", timestamp=old + 5))

    # Topic 3: memory audit
    log.append(ConversationAppendRequest(session_id="s1", role="user", content="Run precision audit on cortex", timestamp=old + 6))
    log.append(ConversationAppendRequest(session_id="s1", role="assistant", content="30 queries: GOOD=8, PARTIAL=11, BAD=11", timestamp=old + 7))
    log.append(ConversationAppendRequest(session_id="s1", role="user", content="what was the weighted score?", timestamp=old + 8))
    log.append(ConversationAppendRequest(session_id="s1", role="assistant", content="Weighted: 45%", timestamp=old + 9))

    return log


class TestPrepareTurns:
    def test_each_entry_is_separate_turn(self):
        """Each user/assistant entry becomes its own turn."""
        d = EpisodeDigestor(log=None, cortex=AsyncMock())
        entries = [
            make_entry(id=1, role="user", content="Fix CORS"),
            make_entry(id=2, role="assistant", content="Added UseCors()"),
            make_entry(id=3, role="user", content="test it"),
            make_entry(id=4, role="assistant", content="All green"),
        ]
        turns = d._prepare_turns(entries)
        assert len(turns) == 4
        assert "Fix CORS" in turns[0]["full_text"]
        assert "USER:" in turns[0]["full_text"]
        assert turns[0]["entry_ids"] == [1]
        assert "UseCors" in turns[1]["full_text"]
        assert "ASSISTANT:" in turns[1]["full_text"]
        assert turns[1]["entry_ids"] == [2]

    def test_keeps_short_entries(self):
        """Even short entries are kept — no length filtering."""
        d = EpisodeDigestor(log=None, cortex=AsyncMock())
        entries = [
            make_entry(id=1, role="user", content="ok"),
            make_entry(id=2, role="assistant", content="sure"),
            make_entry(id=3, role="user", content="This is a longer question"),
        ]
        turns = d._prepare_turns(entries)
        assert len(turns) == 3

    def test_skips_tool_use_and_tool_result(self):
        """tool_use and tool_result roles are filtered out."""
        d = EpisodeDigestor(log=None, cortex=AsyncMock())
        entries = [
            make_entry(id=1, role="user", content="read the file and analyze contents"),
            make_entry(id=2, role="tool_use", content='{"tool": "Read", "path": "/foo"}'),
            make_entry(id=3, role="tool_result", content="file contents here blah blah"),
            make_entry(id=4, role="assistant", content="The file contains a Foo class implementation"),
        ]
        turns = d._prepare_turns(entries)
        assert len(turns) == 2
        assert turns[0]["entry_ids"] == [1]
        assert turns[1]["entry_ids"] == [4]

    def test_strips_xml_but_keeps_thinking(self):
        """XML tags stripped, but thinking content preserved as text."""
        d = EpisodeDigestor(log=None, cortex=AsyncMock())
        entries = [
            make_entry(id=1, role="assistant", content="<thinking>internal reasoning about the bug</thinking>The answer is correct."),
        ]
        turns = d._prepare_turns(entries)
        assert len(turns) == 1
        # XML tags stripped but content kept
        assert "internal reasoning" in turns[0]["full_text"]
        assert "answer is correct" in turns[0]["full_text"]
        assert "<thinking>" not in turns[0]["full_text"]

    def test_empty_entries(self):
        d = EpisodeDigestor(log=None, cortex=AsyncMock())
        assert d._prepare_turns([]) == []


class TestTopicSegmentation:
    @pytest.mark.asyncio
    async def test_segments_multiple_topics(self, tmp_path):
        """Classifier segments turns into topic clusters."""
        log = _make_multi_topic_log(tmp_path)
        cortex = AsyncMock()
        cortex.ingest_episode = AsyncMock(return_value="ep_123")
        classifier = AsyncMock()
        classifier.segment_topics = AsyncMock(return_value=[
            {"start": 0, "end": 3, "topic": "CORS fix"},
            {"start": 4, "end": 5, "topic": "systemd deploy"},
            {"start": 6, "end": 9, "topic": "precision audit"},
        ])

        d = EpisodeDigestor(log=log, cortex=cortex, classifier=classifier)
        await d._process_all()

        assert classifier.segment_topics.call_count == 1
        assert cortex.ingest_episode.call_count == 3
        assert log.get_undigested() == []

    @pytest.mark.asyncio
    async def test_processes_all_segments(self, tmp_path):
        """All segments processed, nothing skipped."""
        db = str(tmp_path / "log.db")
        log = ConversationLog(db)
        now = time.time()

        log.append(ConversationAppendRequest(session_id="s1", role="user", content="old topic", timestamp=now - 120))
        log.append(ConversationAppendRequest(session_id="s1", role="assistant", content="old answer", timestamp=now - 119))
        log.append(ConversationAppendRequest(session_id="s1", role="user", content="new topic", timestamp=now - 5))
        log.append(ConversationAppendRequest(session_id="s1", role="assistant", content="new answer", timestamp=now - 4))

        cortex = AsyncMock()
        cortex.ingest_episode = AsyncMock(return_value="ep_ok")
        classifier = AsyncMock()
        classifier.segment_topics = AsyncMock(return_value=[
            {"start": 0, "end": 1, "topic": "old stuff"},
            {"start": 2, "end": 3, "topic": "new stuff"},
        ])

        d = EpisodeDigestor(log=log, cortex=cortex, classifier=classifier)
        result = await d._process_all()

        assert result is True
        assert cortex.ingest_episode.call_count == 2
        assert len(log.get_undigested()) == 0


class TestProcessAll:
    @pytest.mark.asyncio
    async def test_noop_when_empty(self, tmp_path):
        """_process_all() returns False with no undigested entries."""
        db = str(tmp_path / "log.db")
        log = ConversationLog(db)
        cortex = AsyncMock()
        d = EpisodeDigestor(log=log, cortex=cortex)
        result = await d._process_all()
        assert result is False
        assert cortex.ingest_episode.call_count == 0

    @pytest.mark.asyncio
    async def test_no_classifier_single_segment(self, tmp_path):
        """Without classifier, all turns become 1 segment."""
        log = _make_multi_topic_log(tmp_path)
        cortex = AsyncMock()
        cortex.ingest_episode = AsyncMock(return_value="ep_all")
        d = EpisodeDigestor(log=log, cortex=cortex, classifier=None)
        await d._process_all()

        assert cortex.ingest_episode.call_count == 1
        assert log.get_undigested() == []

    @pytest.mark.asyncio
    async def test_double_process_no_reprocess(self, tmp_path):
        """Second call doesn't re-process already digested entries."""
        log = _make_multi_topic_log(tmp_path)
        cortex = AsyncMock()
        cortex.ingest_episode = AsyncMock(return_value="ep_x")
        classifier = AsyncMock()
        classifier.segment_topics = AsyncMock(return_value=[
            {"start": 0, "end": 9, "topic": "everything"},
        ])
        d = EpisodeDigestor(log=log, cortex=cortex, classifier=classifier)
        await d._process_all()
        await d._process_all()
        assert cortex.ingest_episode.call_count == 1

    @pytest.mark.asyncio
    async def test_no_log(self):
        """_process_all() does nothing with no log."""
        d = EpisodeDigestor(log=None, cortex=AsyncMock())
        result = await d._process_all()
        assert result is False

    @pytest.mark.asyncio
    async def test_classifier_fallback_on_error(self, tmp_path):
        """If classifier fails, fallback to chunks of 5."""
        log = _make_multi_topic_log(tmp_path)
        cortex = AsyncMock()
        cortex.ingest_episode = AsyncMock(return_value="ep_x")
        classifier = AsyncMock()
        classifier.segment_topics = AsyncMock(side_effect=Exception("LLM down"))

        d = EpisodeDigestor(log=log, cortex=cortex, classifier=classifier)
        await d._process_all()

        assert cortex.ingest_episode.call_count >= 1
        assert log.get_undigested() == []

    @pytest.mark.asyncio
    async def test_carryover_on_multi_batch(self, tmp_path):
        """When data exceeds token budget, last segment carries over to next batch."""
        orig_budget = cortex_config.DIGESTOR_BUDGET

        # Set tiny budget so each batch fits ~2 turns
        cortex_config.DIGESTOR_BUDGET = 8
        try:
            log = _make_multi_topic_log(tmp_path)
            cortex = AsyncMock()
            cortex.ingest_episode = AsyncMock(return_value="ep_x")
            classifier = AsyncMock()
            classifier.segment_topics = AsyncMock(return_value=[
                {"start": 0, "end": 0, "topic": "A"},
                {"start": 1, "end": 1, "topic": "B"},
            ])

            d = EpisodeDigestor(log=log, cortex=cortex, classifier=classifier)
            result = await d._process_all()

            assert result is True
            assert log.get_undigested() == []
            assert cortex.ingest_episode.call_count >= 2
        finally:
            cortex_config.DIGESTOR_BUDGET = orig_budget


class TestDistillMode:
    @pytest.mark.asyncio
    async def test_process_all_with_distiller(self, tmp_path):
        """With distiller, calls ingest_knowledge_fast for each entry."""
        log = _make_multi_topic_log(tmp_path)
        cortex = AsyncMock()
        cortex.ingest_knowledge_fast = AsyncMock(return_value="kn_123")
        distiller = AsyncMock()
        distiller.distill = AsyncMock(return_value=[
            KnowledgeEntry(topic="CORS fix", content="Added UseCors().", type="bugfix", project="neo-reloaded"),
            KnowledgeEntry(topic="Systemd deploy", content="Created service file.", type="config", project="neo-reloaded"),
        ])

        d = EpisodeDigestor(log=log, cortex=cortex, distiller=distiller)
        result = await d._process_all()

        assert result is True
        assert distiller.distill.call_count == 1
        assert cortex.ingest_knowledge_fast.call_count == 2
        assert log.get_undigested() == []

    @pytest.mark.asyncio
    async def test_distill_uses_fast_ingest(self, tmp_path):
        """Distill mode calls ingest_knowledge_fast, NOT ingest_knowledge."""
        log = _make_multi_topic_log(tmp_path)
        cortex = AsyncMock()
        cortex.ingest_knowledge_fast = AsyncMock(return_value="kn_fast")
        cortex.ingest_knowledge = AsyncMock()  # should NOT be called
        distiller = AsyncMock()
        distiller.distill = AsyncMock(return_value=[
            KnowledgeEntry(topic="Test", content="Content.", type="feature", project="test"),
        ])

        d = EpisodeDigestor(log=log, cortex=cortex, distiller=distiller)
        await d._process_all()

        cortex.ingest_knowledge_fast.assert_called()
        cortex.ingest_knowledge.assert_not_called()

    @pytest.mark.asyncio
    async def test_distill_fallback_on_error(self, tmp_path):
        """If distiller fails, entries still marked digested."""
        log = _make_multi_topic_log(tmp_path)
        cortex = AsyncMock()
        distiller = AsyncMock()
        distiller.distill = AsyncMock(side_effect=Exception("Gemini down"))

        d = EpisodeDigestor(log=log, cortex=cortex, distiller=distiller)
        result = await d._process_all()

        assert result is True
        assert log.get_undigested() == []

    @pytest.mark.asyncio
    async def test_distill_empty_result(self, tmp_path):
        """Distiller returns empty list → no ingest calls, but entries digested."""
        log = _make_multi_topic_log(tmp_path)
        cortex = AsyncMock()
        cortex.ingest_knowledge_fast = AsyncMock()
        distiller = AsyncMock()
        distiller.distill = AsyncMock(return_value=[])

        d = EpisodeDigestor(log=log, cortex=cortex, distiller=distiller)
        await d._process_all()

        cortex.ingest_knowledge_fast.assert_not_called()
        assert log.get_undigested() == []

    @pytest.mark.asyncio
    async def test_fallback_to_segment_without_distiller(self, tmp_path):
        """Without distiller, falls back to segment mode (ingest_episode)."""
        log = _make_multi_topic_log(tmp_path)
        cortex = AsyncMock()
        cortex.ingest_episode = AsyncMock(return_value="ep_all")

        d = EpisodeDigestor(log=log, cortex=cortex, classifier=None, distiller=None)
        await d._process_all()

        assert cortex.ingest_episode.call_count == 1
        assert log.get_undigested() == []


class TestDigestorConsolidation:
    """Test consolidation trigger on idle."""

    def test_consolidator_accepted(self):
        """Digestor accepts consolidator parameter."""
        consolidator = AsyncMock()
        d = EpisodeDigestor(log=None, cortex=AsyncMock(), consolidator=consolidator)
        assert d._consolidator is consolidator

    def test_consolidator_optional(self):
        """Digestor works without consolidator."""
        d = EpisodeDigestor(log=None, cortex=AsyncMock())
        assert d._consolidator is None


class TestDigestorPause:
    """Test pause/resume for memory_rebuild."""

    def test_pause_resume(self):
        d = EpisodeDigestor(log=None, cortex=AsyncMock())
        assert not d._paused
        d.pause()
        assert d._paused
        d.resume()
        assert not d._paused
