"""U6: Lead planning-time historical reflection retrieval (R12, R13).

Covers:
- ReflexionEngine.retrieve_prompt_reflection(): returns best reflection by score
- Score filtering: score <= 0.5 not returned
- No episodic_memory → returns None
- Retrieval failure → returns None (non-blocking)
- TeamOrchestrator._decompose_task: prepends improved_prompt when reflection found
- No reflexion_engine → default prompt (backward compat)
- Retrieval failure → default prompt (non-blocking)
"""

from __future__ import annotations

from unittest.mock import AsyncMock, MagicMock

import pytest

from agentkit.core.reflexion import ReflexionEngine
from agentkit.experts.orchestrator import TeamOrchestrator
from agentkit.experts.team import ExpertTeam
from agentkit.memory.base import MemoryItem

from tests.unit.experts._helpers import make_execute_stream_mock


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


def _make_memory_item(
    score: float = 0.8,
    output_summary: str = "improved prompt text",
    reflection: str = "reflection text",
    version: int = 1,
) -> MemoryItem:
    from datetime import datetime, timezone

    # Shape matches EpisodicMemory.list_prompt_reflections_by_hash return:
    # output_summary/reflection/quality_score live in value dict,
    # version/task_hash live in metadata.
    return MemoryItem(
        key="prompt_reflection:abc:1",
        value={
            "input_summary": "test",
            "output_summary": output_summary,
            "reflection": reflection,
            "quality_score": score,
        },
        metadata={
            "task_type": "prompt_reflection",
            "version": version,
            "task_hash": "abc",
            "created_at": datetime.now(timezone.utc).isoformat(),
        },
        score=score,
        created_at=datetime.now(timezone.utc),
    )


def _make_llm_gateway_mock() -> MagicMock:
    gw = MagicMock()
    response = MagicMock()
    response.content = (
        '[{"name":"A","assigned_expert":"lead","task_description":"a","depends_on":[]}]'
    )
    gw.chat = AsyncMock(return_value=response)
    return gw


def _make_team_with_experts() -> ExpertTeam:
    from agentkit.core.handoff_transport import InProcessHandoffTransport
    from agentkit.core.protocol import TaskResult, TaskStatus
    from agentkit.experts.config import ExpertConfig
    from agentkit.experts.expert import Expert

    team = ExpertTeam()
    team._handoff_transport = AsyncMock(spec=InProcessHandoffTransport)

    config = ExpertConfig(
        name="lead",
        agent_type="expert",
        persona="测试",
        thinking_style="逻辑",
        bound_skills=["s"],
        is_lead=True,
        task_mode="llm_generate",
        prompt={"identity": "测试"},
    )
    expert = MagicMock(spec=Expert)
    expert.config = config
    expert.is_active = True
    expert.team_id = None
    expert.get_capabilities_summary.return_value = {"name": "lead"}

    mock_agent = MagicMock()
    mock_agent.execute = AsyncMock(
        return_value=TaskResult(
            task_id="t",
            agent_name="lead",
            status=TaskStatus.COMPLETED.value,
            output_data={"content": "result"},
            error_message=None,
            started_at=None,
            completed_at=None,
        )
    )
    mock_agent.execute_stream = make_execute_stream_mock("result")
    mock_agent._llm_gateway = None
    expert.agent = mock_agent

    team._experts["lead"] = expert
    team._lead_expert_name = "lead"
    return team


# ── ReflexionEngine.retrieve_prompt_reflection ─────────────────────────


class TestRetrievePromptReflection:
    """U6/R12: retrieve_prompt_reflection returns best reflection by score."""

    @pytest.mark.asyncio
    async def test_returns_none_when_no_episodic_memory(self):
        gw = MagicMock()
        engine = ReflexionEngine(llm_gateway=gw, episodic_memory=None)
        result = await engine.retrieve_prompt_reflection(task_input="test")
        assert result is None

    @pytest.mark.asyncio
    async def test_returns_best_reflection_by_score(self):
        episodic = MagicMock()
        # Two results: score 0.6 and 0.9 — should return 0.9
        items = [
            _make_memory_item(score=0.6, output_summary="prompt v1"),
            _make_memory_item(score=0.9, output_summary="prompt v2"),
        ]
        episodic.search_prompt_reflections = AsyncMock(return_value=items)
        gw = MagicMock()
        engine = ReflexionEngine(llm_gateway=gw, episodic_memory=episodic)

        result = await engine.retrieve_prompt_reflection(task_input="test")

        assert result is not None
        assert result["score"] == 0.9
        assert result["improved_prompt"] == "prompt v2"

    @pytest.mark.asyncio
    async def test_filters_low_score_reflections(self):
        """score <= 0.5 should not be returned."""
        episodic = MagicMock()
        items = [_make_memory_item(score=0.3, output_summary="low score")]
        episodic.search_prompt_reflections = AsyncMock(return_value=items)
        gw = MagicMock()
        engine = ReflexionEngine(llm_gateway=gw, episodic_memory=episodic)

        result = await engine.retrieve_prompt_reflection(task_input="test", min_score=0.5)
        assert result is None

    @pytest.mark.asyncio
    async def test_returns_none_when_no_results(self):
        episodic = MagicMock()
        episodic.search_prompt_reflections = AsyncMock(return_value=[])
        gw = MagicMock()
        engine = ReflexionEngine(llm_gateway=gw, episodic_memory=episodic)

        result = await engine.retrieve_prompt_reflection(task_input="test")
        assert result is None

    @pytest.mark.asyncio
    async def test_returns_none_on_search_failure(self):
        episodic = MagicMock()
        episodic.search_prompt_reflections = AsyncMock(side_effect=RuntimeError("DB down"))
        gw = MagicMock()
        engine = ReflexionEngine(llm_gateway=gw, episodic_memory=episodic)

        result = await engine.retrieve_prompt_reflection(task_input="test")
        assert result is None

    @pytest.mark.asyncio
    async def test_returns_reflection_fields_complete(self):
        episodic = MagicMock()
        items = [
            _make_memory_item(
                score=0.85,
                output_summary="improved prompt",
                reflection="what went wrong",
                version=3,
            )
        ]
        episodic.search_prompt_reflections = AsyncMock(return_value=items)
        gw = MagicMock()
        engine = ReflexionEngine(llm_gateway=gw, episodic_memory=episodic)

        result = await engine.retrieve_prompt_reflection(task_input="test")

        assert result is not None
        assert result["improved_prompt"] == "improved prompt"
        assert result["reflection"] == "what went wrong"
        assert result["version"] == 3
        assert result["score"] == 0.85
        assert "task_hash" in result

    @pytest.mark.asyncio
    async def test_custom_min_score_threshold(self):
        """min_score=0.7 filters out score=0.6."""
        episodic = MagicMock()
        items = [_make_memory_item(score=0.6, output_summary="medium")]
        episodic.search_prompt_reflections = AsyncMock(return_value=items)
        gw = MagicMock()
        engine = ReflexionEngine(llm_gateway=gw, episodic_memory=episodic)

        result = await engine.retrieve_prompt_reflection(task_input="test", min_score=0.7)
        assert result is None


# ── TeamOrchestrator._decompose_task with reflection ──────────────────


class TestDecomposeWithReflection:
    """U6/R13: _decompose_task prepends improved_prompt when reflection found."""

    @pytest.mark.asyncio
    async def test_decompose_prepends_reflection_when_found(self):
        """When reflexion_engine returns a reflection, the decomposition
        prompt includes the improved_prompt."""
        team = _make_team_with_experts()
        gw = _make_llm_gateway_mock()
        team._experts["lead"].agent._llm_gateway = gw

        # Mock reflexion_engine
        reflexion = MagicMock(spec=ReflexionEngine)
        reflexion.retrieve_prompt_reflection = AsyncMock(
            return_value={
                "improved_prompt": "USE THIS IMPROVED APPROACH",
                "score": 0.85,
                "reflection": "past mistake",
                "version": 2,
                "task_hash": "abc",
            }
        )

        orchestrator = TeamOrchestrator(team, reflexion_engine=reflexion)
        await orchestrator._decompose_task(
            team.lead_expert, "test task", trigger_reason="verify_retry"
        )

        # Verify retrieve was called
        reflexion.retrieve_prompt_reflection.assert_awaited_once()
        # Verify the prompt sent to LLM includes the improved_prompt
        call_kwargs = gw.chat.await_args.kwargs
        messages = call_kwargs.get("messages") or gw.chat.await_args.args[0]
        prompt_content = messages[0]["content"] if isinstance(messages, list) else str(messages)
        assert "USE THIS IMPROVED APPROACH" in prompt_content
        assert "Historical Reflection" in prompt_content

    @pytest.mark.asyncio
    async def test_decompose_uses_default_when_no_reflexion_engine(self):
        """No reflexion_engine → default prompt (backward compat)."""
        team = _make_team_with_experts()
        gw = _make_llm_gateway_mock()
        team._experts["lead"].agent._llm_gateway = gw

        orchestrator = TeamOrchestrator(team, reflexion_engine=None)
        await orchestrator._decompose_task(team.lead_expert, "test task")

        # Verify default prompt (no Historical Reflection section)
        call_kwargs = gw.chat.await_args.kwargs
        messages = call_kwargs.get("messages") or gw.chat.await_args.args[0]
        prompt_content = messages[0]["content"] if isinstance(messages, list) else str(messages)
        assert "Historical Reflection" not in prompt_content

    @pytest.mark.asyncio
    async def test_decompose_uses_default_when_no_reflection_found(self):
        """reflexion_engine returns None → default prompt."""
        team = _make_team_with_experts()
        gw = _make_llm_gateway_mock()
        team._experts["lead"].agent._llm_gateway = gw

        reflexion = MagicMock(spec=ReflexionEngine)
        reflexion.retrieve_prompt_reflection = AsyncMock(return_value=None)

        orchestrator = TeamOrchestrator(team, reflexion_engine=reflexion)
        await orchestrator._decompose_task(
            team.lead_expert, "test task", trigger_reason="verify_retry"
        )

        call_kwargs = gw.chat.await_args.kwargs
        messages = call_kwargs.get("messages") or gw.chat.await_args.args[0]
        prompt_content = messages[0]["content"] if isinstance(messages, list) else str(messages)
        assert "Historical Reflection" not in prompt_content

    @pytest.mark.asyncio
    async def test_decompose_uses_default_when_retrieval_fails(self):
        """reflexion_engine.retrieve raises → default prompt (non-blocking)."""
        team = _make_team_with_experts()
        gw = _make_llm_gateway_mock()
        team._experts["lead"].agent._llm_gateway = gw

        reflexion = MagicMock(spec=ReflexionEngine)
        reflexion.retrieve_prompt_reflection = AsyncMock(side_effect=RuntimeError("search failed"))

        orchestrator = TeamOrchestrator(team, reflexion_engine=reflexion)
        await orchestrator._decompose_task(
            team.lead_expert, "test task", trigger_reason="verify_retry"
        )

        # Default prompt used despite retrieval failure
        call_kwargs = gw.chat.await_args.kwargs
        messages = call_kwargs.get("messages") or gw.chat.await_args.args[0]
        prompt_content = messages[0]["content"] if isinstance(messages, list) else str(messages)
        assert "Historical Reflection" not in prompt_content

    @pytest.mark.asyncio
    async def test_decompose_skips_reflection_when_no_improved_prompt(self):
        """reflexion_engine returns dict without improved_prompt → no hint."""
        team = _make_team_with_experts()
        gw = _make_llm_gateway_mock()
        team._experts["lead"].agent._llm_gateway = gw

        reflexion = MagicMock(spec=ReflexionEngine)
        reflexion.retrieve_prompt_reflection = AsyncMock(
            return_value={"improved_prompt": "", "score": 0.8}  # empty improved_prompt
        )

        orchestrator = TeamOrchestrator(team, reflexion_engine=reflexion)
        await orchestrator._decompose_task(
            team.lead_expert, "test task", trigger_reason="verify_retry"
        )

        call_kwargs = gw.chat.await_args.kwargs
        messages = call_kwargs.get("messages") or gw.chat.await_args.args[0]
        prompt_content = messages[0]["content"] if isinstance(messages, list) else str(messages)
        assert "Historical Reflection" not in prompt_content


# ── U4/KTD5: retrieve_prompt_reflection trigger gating ────────────────


class TestDecomposeTriggerGating:
    """U4/KTD5: _decompose_task only calls retrieve_prompt_reflection when
    trigger_reason is one of {verify_retry, schema_validation, loop_detection}.
    Non-trigger decompositions skip retrieval entirely (default prompt)."""

    def _make_orchestrator_with_reflexion(
        self, retrieve_return: object = None, retrieve_side_effect: object | None = None
    ) -> tuple[TeamOrchestrator, MagicMock, MagicMock]:
        """Build orchestrator + gateway + reflexion mock.

        Returns (orchestrator, gateway, reflexion). retrieve_side_effect, if
        set, takes precedence over retrieve_return.
        """
        team = _make_team_with_experts()
        gw = _make_llm_gateway_mock()
        team._experts["lead"].agent._llm_gateway = gw

        reflexion = MagicMock(spec=ReflexionEngine)
        if retrieve_side_effect is not None:
            reflexion.retrieve_prompt_reflection = AsyncMock(side_effect=retrieve_side_effect)
        else:
            reflexion.retrieve_prompt_reflection = AsyncMock(return_value=retrieve_return)

        orchestrator = TeamOrchestrator(team, reflexion_engine=reflexion)
        return orchestrator, gw, reflexion

    @pytest.mark.asyncio
    async def test_retrieves_on_verify_retry_trigger(self):
        orchestrator, _, reflexion = self._make_orchestrator_with_reflexion(
            retrieve_return={
                "improved_prompt": "IMPROVED",
                "score": 0.8,
                "reflection": "r",
                "version": 1,
                "task_hash": "abc",
            }
        )
        await orchestrator._decompose_task(
            orchestrator._team.lead_expert, "task", trigger_reason="verify_retry"
        )
        reflexion.retrieve_prompt_reflection.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_retrieves_on_schema_validation_trigger(self):
        orchestrator, _, reflexion = self._make_orchestrator_with_reflexion(retrieve_return=None)
        await orchestrator._decompose_task(
            orchestrator._team.lead_expert, "task", trigger_reason="schema_validation"
        )
        reflexion.retrieve_prompt_reflection.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_retrieves_on_loop_detection_trigger(self):
        orchestrator, _, reflexion = self._make_orchestrator_with_reflexion(retrieve_return=None)
        await orchestrator._decompose_task(
            orchestrator._team.lead_expert, "task", trigger_reason="loop_detection"
        )
        reflexion.retrieve_prompt_reflection.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_skips_retrieval_when_trigger_is_none(self):
        """Fresh decomposition (trigger_reason=None) → no retrieval."""
        orchestrator, _, reflexion = self._make_orchestrator_with_reflexion(
            retrieve_return={"improved_prompt": "would-be-used", "score": 0.9}
        )
        await orchestrator._decompose_task(orchestrator._team.lead_expert, "task")
        reflexion.retrieve_prompt_reflection.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_skips_retrieval_on_invalid_trigger_reason(self):
        """trigger_reason not in the allowed set → no retrieval."""
        orchestrator, _, reflexion = self._make_orchestrator_with_reflexion(
            retrieve_return={"improved_prompt": "would-be-used", "score": 0.9}
        )
        await orchestrator._decompose_task(
            orchestrator._team.lead_expert, "task", trigger_reason="random"
        )
        reflexion.retrieve_prompt_reflection.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_triggered_retrieval_failure_degrades_to_default_prompt(self):
        """trigger_reason valid but retrieve raises → default prompt (non-blocking)."""
        orchestrator, gw, reflexion = self._make_orchestrator_with_reflexion(
            retrieve_side_effect=RuntimeError("search down")
        )
        await orchestrator._decompose_task(
            orchestrator._team.lead_expert, "task", trigger_reason="verify_retry"
        )
        # Retrieval was attempted (triggered) but failed — default prompt used.
        reflexion.retrieve_prompt_reflection.assert_awaited_once()
        call_kwargs = gw.chat.await_args.kwargs
        messages = call_kwargs.get("messages") or gw.chat.await_args.args[0]
        prompt_content = messages[0]["content"] if isinstance(messages, list) else str(messages)
        assert "Historical Reflection" not in prompt_content

    @pytest.mark.asyncio
    async def test_skips_retrieval_when_no_reflexion_engine_even_with_trigger(self):
        """No reflexion_engine → retrieval skipped regardless of trigger_reason."""
        team = _make_team_with_experts()
        gw = _make_llm_gateway_mock()
        team._experts["lead"].agent._llm_gateway = gw

        orchestrator = TeamOrchestrator(team, reflexion_engine=None)
        await orchestrator._decompose_task(
            orchestrator._team.lead_expert, "task", trigger_reason="verify_retry"
        )
        # Default prompt — no Historical Reflection section.
        call_kwargs = gw.chat.await_args.kwargs
        messages = call_kwargs.get("messages") or gw.chat.await_args.args[0]
        prompt_content = messages[0]["content"] if isinstance(messages, list) else str(messages)
        assert "Historical Reflection" not in prompt_content