"""Tests for ReAct engine compression integration (U3)"""

from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from agentkit.core.compressor import CompressionStrategy, ContextCompressor
from agentkit.core.react import ReActEngine
from agentkit.llm.protocol import LLMResponse, TokenUsage, ToolCall


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


def make_mock_gateway() -> MagicMock:
    """创建一个 mock LLMGateway"""
    from agentkit.llm.gateway import LLMGateway

    gateway = MagicMock(spec=LLMGateway)
    response = LLMResponse(
        content="Final answer",
        model="test-model",
        usage=TokenUsage(prompt_tokens=10, completion_tokens=10),
    )
    gateway.chat = AsyncMock(return_value=response)
    return gateway


def make_mock_gateway_with_tool_call() -> MagicMock:
    """创建一个返回 tool_call 的 mock LLMGateway，第二次调用返回最终答案"""
    from agentkit.llm.gateway import LLMGateway

    gateway = MagicMock(spec=LLMGateway)

    # 第一次调用返回 tool_call，第二次返回最终答案
    tool_response = LLMResponse(
        content="",
        model="test-model",
        usage=TokenUsage(prompt_tokens=10, completion_tokens=10),
        tool_calls=[
            ToolCall(id="call_1", name="search", arguments={"query": "test"}),
        ],
    )
    final_response = LLMResponse(
        content="Final answer after tool",
        model="test-model",
        usage=TokenUsage(prompt_tokens=10, completion_tokens=10),
    )
    gateway.chat = AsyncMock(side_effect=[tool_response, final_response])
    return gateway


def make_long_messages(count: int = 10, content_length: int = 2000) -> list[dict]:
    """生成长消息列表用于测试压缩"""
    messages = [{"role": "system", "content": "You are a helpful assistant."}]
    for i in range(count):
        messages.append({
            "role": "user",
            "content": "x" * content_length + f" message {i}",
        })
        messages.append({
            "role": "assistant",
            "content": "y" * content_length + f" reply {i}",
        })
    return messages


def make_mock_compressor() -> MagicMock:
    """创建一个 mock CompressionStrategy"""
    compressor = MagicMock(spec=CompressionStrategy)
    compressor.compress = AsyncMock(return_value=[{"role": "user", "content": "compressed"}])
    compressor.compress_tool_result = AsyncMock(return_value="compressed tool result")
    compressor.is_available = MagicMock(return_value=True)
    return compressor


# ── TestBuildToolResultMessage ────────────────────────


class TestBuildToolResultMessage:
    """_build_tool_result_message 方法测试"""

    async def test_no_compressor_returns_original(self):
        engine = ReActEngine(llm_gateway=make_mock_gateway())
        result = await engine._build_tool_result_message("tc_1", {"key": "value"})
        assert result == {
            "role": "tool",
            "tool_call_id": "tc_1",
            "content": "{'key': 'value'}",
        }

    async def test_with_compressor_calls_compress_tool_result(self):
        compressor = make_mock_compressor()
        engine = ReActEngine(llm_gateway=make_mock_gateway())
        result = await engine._build_tool_result_message(
            "tc_1", {"key": "value"}, compressor=compressor, tool_name="search"
        )
        compressor.compress_tool_result.assert_called_once_with("search", {"key": "value"})
        assert result["content"] == "compressed tool result"
        assert result["role"] == "tool"
        assert result["tool_call_id"] == "tc_1"

    async def test_compressor_failure_falls_back(self):
        compressor = make_mock_compressor()
        compressor.compress_tool_result = AsyncMock(side_effect=RuntimeError("compression error"))
        engine = ReActEngine(llm_gateway=make_mock_gateway())
        result = await engine._build_tool_result_message(
            "tc_1", {"key": "value"}, compressor=compressor, tool_name="search"
        )
        # 应该回退到 str(result)
        assert result["content"] == "{'key': 'value'}"
        assert result["role"] == "tool"

    async def test_compressor_receives_tool_name(self):
        compressor = make_mock_compressor()
        engine = ReActEngine(llm_gateway=make_mock_gateway())
        await engine._build_tool_result_message(
            "tc_1", "some result", compressor=compressor, tool_name="web_crawl"
        )
        compressor.compress_tool_result.assert_called_once_with("web_crawl", "some result")

    async def test_compressor_without_tool_name_skips_compression(self):
        """compressor 存在但 tool_name 为 None 时不压缩"""
        compressor = make_mock_compressor()
        engine = ReActEngine(llm_gateway=make_mock_gateway())
        result = await engine._build_tool_result_message(
            "tc_1", "some result", compressor=compressor, tool_name=None
        )
        compressor.compress_tool_result.assert_not_called()
        assert result["content"] == "some result"


# ── TestShouldCompress ───────────────────────────────


class TestShouldCompress:
    """_should_compress 辅助方法测试"""

    def test_no_compressor_returns_false(self):
        engine = ReActEngine(llm_gateway=make_mock_gateway())
        conversation = [{"role": "user", "content": "x" * 100000}]
        assert engine._should_compress(conversation, None) is False

    def test_short_conversation_returns_false(self):
        compressor = make_mock_compressor()
        engine = ReActEngine(llm_gateway=make_mock_gateway())
        conversation = [
            {"role": "user", "content": "Hello"},
            {"role": "assistant", "content": "Hi there"},
        ]
        assert engine._should_compress(conversation, compressor) is False

    def test_long_conversation_returns_true(self):
        compressor = make_mock_compressor()
        engine = ReActEngine(llm_gateway=make_mock_gateway())
        # 8000 tokens * 4 chars/token = 32000 chars needed
        conversation = [{"role": "user", "content": "x" * 40000}]
        assert engine._should_compress(conversation, compressor) is True

    def test_boundary_at_threshold(self):
        compressor = make_mock_compressor()
        engine = ReActEngine(llm_gateway=make_mock_gateway())
        # Exactly 8000 tokens = 32000 chars → should NOT trigger (> not >=)
        conversation = [{"role": "user", "content": "x" * 32000}]
        assert engine._should_compress(conversation, compressor) is False

    def test_just_above_threshold(self):
        compressor = make_mock_compressor()
        engine = ReActEngine(llm_gateway=make_mock_gateway())
        # 32001 chars → 32001//4 = 8000 tokens, still not > 8000
        # 32004 chars → 32004//4 = 8001 tokens, > 8000
        conversation = [{"role": "user", "content": "x" * 32004}]
        assert engine._should_compress(conversation, compressor) is True


# ── TestReActLoopCompression ─────────────────────────


class TestReActLoopCompression:
    """ReAct 循环内压缩集成测试"""

    async def test_tool_results_compressed_in_loop(self):
        """工具结果在循环中被压缩后拼入 conversation"""
        compressor = make_mock_compressor()
        compressor.compress_tool_result = AsyncMock(return_value="COMPRESSED:search_result")

        gateway = make_mock_gateway_with_tool_call()
        engine = ReActEngine(llm_gateway=gateway)

        # 注册一个 mock tool
        mock_tool = MagicMock()
        mock_tool.name = "search"
        mock_tool.safe_execute = AsyncMock(return_value="original search result")

        result = await engine.execute(
            messages=[{"role": "user", "content": "Search for test"}],
            tools=[mock_tool],
            compressor=compressor,
        )

        # 验证 compress_tool_result 被调用
        compressor.compress_tool_result.assert_called_once_with("search", "original search result")

    async def test_incremental_compression_triggered(self):
        """长对话触发增量压缩 compressor.compress()"""
        compressor = make_mock_compressor()
        # 让 compress 返回压缩后的短对话
        compressor.compress = AsyncMock(return_value=[
            {"role": "system", "content": "Summary"},
            {"role": "user", "content": "recent"},
        ])

        gateway = make_mock_gateway_with_tool_call()
        engine = ReActEngine(llm_gateway=gateway)

        # 构造一个很长的初始消息，使 conversation 超过阈值
        long_content = "x" * 40000
        mock_tool = MagicMock()
        mock_tool.name = "search"
        mock_tool.safe_execute = AsyncMock(return_value="result")

        result = await engine.execute(
            messages=[{"role": "user", "content": long_content}],
            tools=[mock_tool],
            compressor=compressor,
        )

        # 验证增量压缩被触发（compress 被调用了至少一次）
        # 初始 compress 在 L218-222，增量 compress 在工具结果后
        assert compressor.compress.call_count >= 1

    async def test_incremental_compression_failure_handled(self):
        """compressor.compress() 异常时循环继续"""
        compressor = make_mock_compressor()
        # 第一次 compress 调用成功（初始压缩），增量压缩时失败
        call_count = 0

        async def compress_side_effect(messages):
            nonlocal call_count
            call_count += 1
            if call_count > 1:
                raise RuntimeError("Incremental compression failed")
            return messages

        compressor.compress = AsyncMock(side_effect=compress_side_effect)

        gateway = make_mock_gateway_with_tool_call()
        engine = ReActEngine(llm_gateway=gateway)

        # 构造长消息触发增量压缩
        long_content = "x" * 40000
        mock_tool = MagicMock()
        mock_tool.name = "search"
        mock_tool.safe_execute = AsyncMock(return_value="result")

        # 不应该抛出异常
        result = await engine.execute(
            messages=[{"role": "user", "content": long_content}],
            tools=[mock_tool],
            compressor=compressor,
        )

        # 应该正常返回结果
        assert result.output == "Final answer after tool"

    async def test_no_compressor_backward_compatible(self):
        """compressor=None 时行为与之前完全一致"""
        gateway = make_mock_gateway_with_tool_call()
        engine = ReActEngine(llm_gateway=gateway)

        mock_tool = MagicMock()
        mock_tool.name = "search"
        mock_tool.safe_execute = AsyncMock(return_value="search result data")

        result = await engine.execute(
            messages=[{"role": "user", "content": "Search for test"}],
            tools=[mock_tool],
            compressor=None,
        )

        assert result.output == "Final answer after tool"
        assert result.status == "success"
        assert len(result.trajectory) == 2  # 1 tool_call + 1 final_answer

    async def test_execute_stream_with_compressor(self):
        """execute_stream 模式下压缩正常工作"""
        compressor = make_mock_compressor()
        compressor.compress_tool_result = AsyncMock(return_value="COMPRESSED:result")

        gateway = make_mock_gateway_with_tool_call()
        engine = ReActEngine(llm_gateway=gateway)

        mock_tool = MagicMock()
        mock_tool.name = "search"
        mock_tool.safe_execute = AsyncMock(return_value="original result")

        events = []
        async for event in engine.execute_stream(
            messages=[{"role": "user", "content": "Search for test"}],
            tools=[mock_tool],
            compressor=compressor,
        ):
            events.append(event)

        # 验证 compress_tool_result 被调用
        compressor.compress_tool_result.assert_called_once_with("search", "original result")

        # 验证事件流完整
        event_types = [e.event_type for e in events]
        assert "thinking" in event_types
        assert "tool_call" in event_types
        assert "tool_result" in event_types
        assert "final_answer" in event_types

    async def test_execute_stream_incremental_compression(self):
        """execute_stream 模式下增量压缩触发"""
        compressor = make_mock_compressor()
        compressor.compress = AsyncMock(return_value=[
            {"role": "system", "content": "Summary"},
            {"role": "user", "content": "recent"},
        ])

        gateway = make_mock_gateway_with_tool_call()
        engine = ReActEngine(llm_gateway=gateway)

        long_content = "x" * 40000
        mock_tool = MagicMock()
        mock_tool.name = "search"
        mock_tool.safe_execute = AsyncMock(return_value="result")

        events = []
        async for event in engine.execute_stream(
            messages=[{"role": "user", "content": long_content}],
            tools=[mock_tool],
            compressor=compressor,
        ):
            events.append(event)

        # 验证增量压缩被触发
        assert compressor.compress.call_count >= 1

    async def test_context_compressor_satisfies_protocol(self):
        """ContextCompressor 满足 CompressionStrategy Protocol"""
        compressor = ContextCompressor()
        assert isinstance(compressor, CompressionStrategy)

    async def test_context_compressor_compress_tool_result_default(self):
        """ContextCompressor.compress_tool_result 默认返回 str(result)"""
        compressor = ContextCompressor()
        result = await compressor.compress_tool_result("search", {"key": "value"})
        assert result == "{'key': 'value'}"


# ── TestOTelSpanLifecycle (P0 fix: span leak) ──────────────────


class TestOTelSpanLifecycle:
    """测试 OTel span 生命周期 — 异常时 span 必须正确关闭"""

    async def test_span_closed_on_success(self):
        """正常执行时 span 被正确关闭"""
        gateway = make_mock_gateway()
        engine = ReActEngine(llm_gateway=gateway)

        mock_span = MagicMock()
        mock_span_cm = MagicMock()
        mock_span_cm.__enter__ = MagicMock(return_value=mock_span)
        mock_span_cm.__exit__ = MagicMock(return_value=False)

        with patch("agentkit.core.react.start_span", return_value=mock_span_cm), \
             patch("agentkit.core.react._OTEL_AVAILABLE", True):
            await engine.execute(messages=[{"role": "user", "content": "hello"}])

        # __exit__ should have been called
        mock_span_cm.__exit__.assert_called_once()

    async def test_span_closed_on_exception(self):
        """LLM 抛出异常时 span 仍被正确关闭"""
        gateway = make_mock_gateway()
        gateway.chat = AsyncMock(side_effect=RuntimeError("LLM error"))
        engine = ReActEngine(llm_gateway=gateway)

        mock_span = MagicMock()
        mock_span_cm = MagicMock()
        mock_span_cm.__enter__ = MagicMock(return_value=mock_span)
        mock_span_cm.__exit__ = MagicMock(return_value=False)

        with patch("agentkit.core.react.start_span", return_value=mock_span_cm), \
             patch("agentkit.core.react._OTEL_AVAILABLE", True):
            with pytest.raises(RuntimeError, match="LLM error"):
                await engine.execute(messages=[{"role": "user", "content": "hello"}])

        # __exit__ must have been called even though exception was raised
        mock_span_cm.__exit__.assert_called_once()

    async def test_span_attributes_set_on_success(self):
        """正常执行时 span 属性被设置"""
        gateway = make_mock_gateway()
        engine = ReActEngine(llm_gateway=gateway)

        mock_span = MagicMock()
        mock_span_cm = MagicMock()
        mock_span_cm.__enter__ = MagicMock(return_value=mock_span)
        mock_span_cm.__exit__ = MagicMock(return_value=False)

        with patch("agentkit.core.react.start_span", return_value=mock_span_cm), \
             patch("agentkit.core.react._OTEL_AVAILABLE", True):
            await engine.execute(messages=[{"role": "user", "content": "hello"}])

        # Verify span attributes were set
        mock_span.set_attribute.assert_any_call("agent.total_steps", 1)
        mock_span.set_attribute.assert_any_call("agent.total_tokens", 20)
        mock_span.set_attribute.assert_any_call("agent.outcome", "success")

    async def test_no_span_when_otel_unavailable(self):
        """_OTEL_AVAILABLE=False 时不创建 span"""
        gateway = make_mock_gateway()
        engine = ReActEngine(llm_gateway=gateway)

        with patch("agentkit.core.react._OTEL_AVAILABLE", False), \
             patch("agentkit.core.react.start_span") as mock_start_span:
            await engine.execute(messages=[{"role": "user", "content": "hello"}])

        # start_span should not be called when OTel is unavailable
        mock_start_span.assert_not_called()
