"""U5: ReflexionEngine reflection persistence to EpisodicMemory (R11, R15).

Covers:
- store_prompt_reflection() stores via existing store() with task_type discriminator
- search_prompt_reflections() filters by task_type="prompt_reflection"
- cleanup_expired() deletes old records (TTL)
- _reflect() persists reflection when episodic_memory configured (non-blocking)
- _reflect() skips persistence when episodic_memory=None (backward compat)
- Persistence failure does not block _reflect (returns reflection text)
- Multi-version coexistence (same task_hash, different versions)
"""

from __future__ import annotations

from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from agentkit.core.react import ReActResult
from agentkit.core.reflexion import ReflexionEngine
from agentkit.memory.episodic import EpisodicMemory


# ── Helpers ────────────────────────────────────────────────────────────


def _make_episodic_memory_mock() -> MagicMock:
    """Create a mock EpisodicMemory that tracks store/search calls."""
    mem = MagicMock(spec=EpisodicMemory)
    mem.store = AsyncMock(return_value=None)
    mem.search = AsyncMock(return_value=[])
    mem.store_prompt_reflection = AsyncMock(return_value="prompt_reflection:abc:1")
    mem.search_prompt_reflections = AsyncMock(return_value=[])
    mem.cleanup_expired = AsyncMock(return_value=0)
    return mem


def _make_react_result(output: str = "test output", status: str = "completed") -> ReActResult:
    return ReActResult(
        output=output,
        trajectory=[],
        total_steps=1,
        total_tokens=10,
        status=status,
    )


def _make_llm_gateway_mock(reflection_text: str = "reflection text") -> MagicMock:
    """Gateway that returns reflection text from chat()."""
    gw = MagicMock()
    response = MagicMock()
    response.content = reflection_text
    gw.chat = AsyncMock(return_value=response)
    return gw


# ── EpisodicMemory.store_prompt_reflection ─────────────────────────────


class TestStorePromptReflection:
    """U5/R11: store_prompt_reflection persists via direct ORM write with
    task_type='prompt_reflection' discriminator and metadata_ JSONB."""

    def _make_persistable_mock(self) -> MagicMock:
        """Build an EpisodicMemory mock with _embedder + _session_factory
        configured so store_prompt_reflection can run its real code path."""
        mem = MagicMock(spec=EpisodicMemory)
        mem._embedder = None  # skip embedding generation in test
        mock_db = AsyncMock()
        mock_db.add = MagicMock()
        mock_db.commit = AsyncMock()
        mock_db.rollback = AsyncMock()
        mem._session_factory = MagicMock()
        mem._session_factory.return_value.__aenter__ = AsyncMock(return_value=mock_db)
        mem._session_factory.return_value.__aexit__ = AsyncMock(return_value=None)
        mem._episodic_model = MagicMock()
        # Track the ORM entry passed to db.add for assertions
        mem._mock_db = mock_db
        return mem

    @pytest.mark.asyncio
    async def test_store_writes_orm_entry_with_correct_fields(self):
        mem = self._make_persistable_mock()
        await EpisodicMemory.store_prompt_reflection(
            mem,
            task_input="test task",
            reflection="reflection text",
            improved_prompt="improved prompt",
            version=1,
            score=0.5,
            agent_name="test_agent",
        )

        # Verify ORM entry created via model constructor + db.add + commit
        mem._episodic_model.assert_called_once()
        call_kwargs = mem._episodic_model.call_args.kwargs
        assert call_kwargs["task_type"] == "prompt_reflection"
        assert call_kwargs["agent_name"] == "test_agent"
        assert call_kwargs["input_summary"] == "test task"
        assert call_kwargs["output_summary"] == "improved prompt"
        assert call_kwargs["reflection"] == "reflection text"
        assert call_kwargs["quality_score"] == 0.5
        assert call_kwargs["outcome"] == "reflection"
        # metadata_ JSONB carries task_hash + version + score
        assert "task_hash" in call_kwargs["metadata_"]
        assert call_kwargs["metadata_"]["version"] == 1
        assert call_kwargs["metadata_"]["score"] == 0.5
        mem._mock_db.add.assert_called_once()
        mem._mock_db.commit.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_store_returns_key_on_success(self):
        mem = self._make_persistable_mock()
        key = await EpisodicMemory.store_prompt_reflection(
            mem,
            task_input="test",
            reflection="r",
            improved_prompt="p",
            version=2,
        )
        assert key is not None
        assert ":2" in key
        assert key.startswith("prompt_reflection:")

    @pytest.mark.asyncio
    async def test_store_returns_none_on_failure(self):
        from sqlalchemy.exc import DBAPIError

        mem = self._make_persistable_mock()
        mem._mock_db.commit = AsyncMock(
            side_effect=DBAPIError("stmt", params={}, orig=Exception("db down"))
        )
        key = await EpisodicMemory.store_prompt_reflection(
            mem,
            task_input="test",
            reflection="r",
            improved_prompt="p",
        )
        assert key is None
        mem._mock_db.rollback.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_store_uses_provided_task_hash(self):
        mem = self._make_persistable_mock()
        key = await EpisodicMemory.store_prompt_reflection(
            mem,
            task_input="test",
            reflection="r",
            improved_prompt="p",
            task_hash="custom_hash",
        )
        assert "custom_hash" in key

    @pytest.mark.asyncio
    async def test_store_generates_task_hash_from_input(self):
        mem = self._make_persistable_mock()
        key1 = await EpisodicMemory.store_prompt_reflection(
            mem, task_input="same task", reflection="r", improved_prompt="p"
        )
        key2 = await EpisodicMemory.store_prompt_reflection(
            mem, task_input="same task", reflection="r", improved_prompt="p"
        )
        # Same task_input → same hash prefix
        assert key1 == key2


# ── EpisodicMemory.search_prompt_reflections ───────────────────────────


class TestSearchPromptReflections:
    """U5/R11: search_prompt_reflections filters by task_type."""

    @pytest.mark.asyncio
    async def test_search_calls_underlying_search_with_filter(self):
        mem = MagicMock(spec=EpisodicMemory)
        mem.search = AsyncMock(return_value=[])
        await EpisodicMemory.search_prompt_reflections(mem, task_input="find similar", top_k=3)

        mem.search.assert_awaited_once()
        call_kwargs = mem.search.await_args.kwargs
        assert call_kwargs["query"] == "find similar"
        assert call_kwargs["top_k"] == 3
        filters = call_kwargs["filters"]
        assert filters["task_type"] == "prompt_reflection"

    @pytest.mark.asyncio
    async def test_search_includes_agent_name_filter_when_provided(self):
        mem = MagicMock(spec=EpisodicMemory)
        mem.search = AsyncMock(return_value=[])
        await EpisodicMemory.search_prompt_reflections(mem, task_input="q", agent_name="agent_x")

        filters = mem.search.await_args.kwargs["filters"]
        assert filters["agent_name"] == "agent_x"

    @pytest.mark.asyncio
    async def test_search_returns_empty_on_failure(self):
        from sqlalchemy.exc import DBAPIError

        mem = MagicMock(spec=EpisodicMemory)
        mem.search = AsyncMock(side_effect=DBAPIError("stmt", params={}, orig=Exception("err")))
        result = await EpisodicMemory.search_prompt_reflections(mem, task_input="q")
        assert result == []


# ── EpisodicMemory.cleanup_expired ─────────────────────────────────────


class TestCleanupExpired:
    """U5/R15: cleanup_expired deletes records older than max_age_days."""

    @pytest.mark.asyncio
    async def test_cleanup_calls_delete_with_cutoff(self):
        # Patch sqlalchemy.delete to bypass ORM model requirement — we only
        # verify the session.execute/commit calls happen.
        mock_db = AsyncMock()
        mock_result = MagicMock()
        mock_result.rowcount = 5
        mock_db.execute = AsyncMock(return_value=mock_result)
        mock_db.commit = AsyncMock()

        mock_session_factory = MagicMock()
        mock_session_factory.return_value.__aenter__ = AsyncMock(return_value=mock_db)
        mock_session_factory.return_value.__aexit__ = AsyncMock(return_value=None)

        # created_at must support `<` comparison with datetime — configure
        # __lt__ to return a truthy mock so `.where(...)` gets a valid arg.
        mock_model = MagicMock()
        mock_model.created_at = MagicMock()
        mock_model.created_at.__lt__ = MagicMock(return_value=MagicMock())

        mem = EpisodicMemory(
            session_factory=mock_session_factory,
            episodic_model=mock_model,
            embedder=None,
            pgvector_enabled=False,
        )

        # Patch sql_delete to return a chainable mock
        fake_stmt = MagicMock()
        fake_stmt.where = MagicMock(return_value=fake_stmt)
        with patch("sqlalchemy.delete", return_value=fake_stmt):
            deleted = await mem.cleanup_expired(max_age_days=30)

        assert deleted == 5
        mock_db.execute.assert_awaited_once()
        mock_db.commit.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_cleanup_returns_zero_on_failure(self):
        from sqlalchemy.exc import DBAPIError

        mock_db = AsyncMock()
        mock_db.execute = AsyncMock(
            side_effect=DBAPIError("stmt", params={}, orig=Exception("err"))
        )
        mock_db.rollback = AsyncMock()

        mock_session_factory = MagicMock()
        mock_session_factory.return_value.__aenter__ = AsyncMock(return_value=mock_db)
        mock_session_factory.return_value.__aexit__ = AsyncMock(return_value=None)

        mock_model = MagicMock()
        mock_model.created_at = MagicMock()
        mock_model.created_at.__lt__ = MagicMock(return_value=MagicMock())

        mem = EpisodicMemory(
            session_factory=mock_session_factory,
            episodic_model=mock_model,
            embedder=None,
            pgvector_enabled=False,
        )

        fake_stmt = MagicMock()
        fake_stmt.where = MagicMock(return_value=fake_stmt)
        with patch("sqlalchemy.delete", return_value=fake_stmt):
            deleted = await mem.cleanup_expired(max_age_days=30)

        assert deleted == 0
        mock_db.rollback.assert_awaited_once()


# ── ReflexionEngine._reflect persistence ──────────────────────────────


class TestReflectPersistence:
    """U5/R11: _reflect persists reflection when episodic_memory configured."""

    @pytest.mark.asyncio
    async def test_reflect_persists_when_episodic_memory_configured(self):
        episodic = _make_episodic_memory_mock()
        gw = _make_llm_gateway_mock(reflection_text="my reflection")
        engine = ReflexionEngine(
            llm_gateway=gw,
            episodic_memory=episodic,
        )

        result = await engine._reflect(
            react_result=_make_react_result(),
            score=0.3,
            messages=[{"role": "user", "content": "test task"}],
            reflect_model="default",
            agent_name="test_agent",
            task_type="test",
        )

        assert result == "my reflection"
        episodic.store_prompt_reflection.assert_awaited_once()
        call_kwargs = episodic.store_prompt_reflection.await_args.kwargs
        assert call_kwargs["task_input"] == "test task"
        assert call_kwargs["reflection"] == "my reflection"
        assert call_kwargs["score"] == 0.3
        assert call_kwargs["agent_name"] == "test_agent"

    @pytest.mark.asyncio
    async def test_reflect_skips_persistence_when_no_episodic_memory(self):
        gw = _make_llm_gateway_mock(reflection_text="my reflection")
        engine = ReflexionEngine(
            llm_gateway=gw,
            episodic_memory=None,  # No persistence
        )

        result = await engine._reflect(
            react_result=_make_react_result(),
            score=0.3,
            messages=[{"role": "user", "content": "test task"}],
            reflect_model="default",
            agent_name="test_agent",
            task_type="test",
        )

        assert result == "my reflection"
        # No episodic_memory → no store call

    @pytest.mark.asyncio
    async def test_reflect_persistence_failure_does_not_block(self):
        """If store_prompt_reflection raises, _reflect still returns reflection."""
        episodic = MagicMock(spec=EpisodicMemory)
        episodic.store_prompt_reflection = AsyncMock(side_effect=RuntimeError("DB down"))
        gw = _make_llm_gateway_mock(reflection_text="important reflection")
        engine = ReflexionEngine(
            llm_gateway=gw,
            episodic_memory=episodic,
        )

        result = await engine._reflect(
            react_result=_make_react_result(),
            score=0.3,
            messages=[{"role": "user", "content": "test task"}],
            reflect_model="default",
            agent_name="test_agent",
            task_type="test",
        )

        # Reflection still returned despite persistence failure
        assert result == "important reflection"
        episodic.store_prompt_reflection.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_reflect_skips_persistence_when_llm_returns_none(self):
        """If LLM returns empty/None, no persistence attempted."""
        episodic = _make_episodic_memory_mock()
        gw = MagicMock()
        response = MagicMock()
        response.content = None  # Empty reflection
        gw.chat = AsyncMock(return_value=response)
        engine = ReflexionEngine(
            llm_gateway=gw,
            episodic_memory=episodic,
        )

        result = await engine._reflect(
            react_result=_make_react_result(),
            score=0.3,
            messages=[{"role": "user", "content": "test task"}],
            reflect_model="default",
            agent_name="test_agent",
            task_type="test",
        )

        assert result is None
        episodic.store_prompt_reflection.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_reflect_skips_persistence_on_llm_failure(self):
        """If LLM call raises, no persistence attempted."""
        episodic = _make_episodic_memory_mock()
        gw = MagicMock()
        gw.chat = AsyncMock(side_effect=RuntimeError("LLM down"))
        engine = ReflexionEngine(
            llm_gateway=gw,
            episodic_memory=episodic,
        )

        result = await engine._reflect(
            react_result=_make_react_result(),
            score=0.3,
            messages=[{"role": "user", "content": "test task"}],
            reflect_model="default",
            agent_name="test_agent",
            task_type="test",
        )

        assert result is None
        episodic.store_prompt_reflection.assert_not_awaited()


# ── Multi-version coexistence ─────────────────────────────────────────


class TestMultiVersionCoexistence:
    """U5/R11: same task_hash with different versions all stored."""

    @pytest.mark.asyncio
    async def test_multiple_versions_stored_with_incrementing_version(self):
        # Reuse the persistable mock helper from TestStorePromptReflection.
        helper = TestStorePromptReflection()
        mem = helper._make_persistable_mock()

        # Store v1, v2, v3 for same task
        key1 = await EpisodicMemory.store_prompt_reflection(
            mem, task_input="same task", reflection="r1", improved_prompt="p1", version=1
        )
        key2 = await EpisodicMemory.store_prompt_reflection(
            mem, task_input="same task", reflection="r2", improved_prompt="p2", version=2
        )
        key3 = await EpisodicMemory.store_prompt_reflection(
            mem, task_input="same task", reflection="r3", improved_prompt="p3", version=3
        )

        # All keys have same task_hash prefix but different version suffix
        assert key1 != key2 != key3
        assert ":1" in key1 and ":2" in key2 and ":3" in key3

        # All three ORM entries created via db.add
        assert mem._mock_db.add.call_count == 3
        assert mem._mock_db.commit.await_count == 3