"""
Unit tests for compaction.py (prune_and_compress_state).

覆盖两部分行为：
- Prune：只动旧 ToolMessage 的 content/metadata；
- Compress：在上下文超阈值时插入 summary 消息（包含 LLM 分支与机械摘要分支）。
"""

from __future__ import annotations

from typing import Any, Dict, List

import pytest
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage

from .compaction import (
    PRUNE_MINIMUM_TOKENS,
    PRUNE_PROTECT_TOKENS,
    PRUNE_PROTECTED_TOOLS,
    prune_and_compress_state,
)


def _make_tool_message(
    content: str,
    tool_call_id: str,
    name: str | None = None,
    metadata: Dict[str, Any] | None = None,
) -> ToolMessage:
    msg = ToolMessage(content=content, tool_call_id=tool_call_id)
    if name is not None:
        # 尽量兼容不同版本的 ToolMessage 实现
        setattr(msg, "name", name)
    if metadata is not None:
        setattr(msg, "metadata", metadata)
    return msg


class TestPruneLogic:
    """测试 prune_and_compress_state 的裁剪行为（不依赖 Compress）。"""

    def test_prune_compacts_only_old_tool_messages_with_enough_tokens(
        self,
        monkeypatch: pytest.MonkeyPatch,
    ):
        """老的 ToolMessage 被 compact，最近两轮用户对话及其工具调用不受影响。"""

        import langchain_agentx.opencode_style_agent.compaction as comp

        # 调整阈值到很小，方便构造测试数据，并避免 Compress 触发
        monkeypatch.setattr(comp, "PRUNE_MINIMUM_TOKENS", 10, raising=False)
        monkeypatch.setattr(comp, "PRUNE_PROTECT_TOKENS", 20, raising=False)
        monkeypatch.setenv("OPENCODE_MAX_CONTEXT_TOKENS", "1000000")
        monkeypatch.setenv("OPENCODE_COMPRESS_TRIGGER_FRACTION", "0.99")

        # 消息序列（从旧到新）：
        #  old_user → old_ai → old_tool(长度200) → recent_user → recent_ai → recent_tool(长度5)
        old_tool = _make_tool_message(content="x" * 200, tool_call_id="old_call", name="normal_tool")
        recent_tool = _make_tool_message(content="y" * 5, tool_call_id="recent_call", name="normal_tool")

        messages: List[Any] = [
            HumanMessage(content="old user 1"),
            AIMessage(content="old ai 1"),
            old_tool,
            HumanMessage(content="recent user 2"),
            AIMessage(content="recent ai 2"),
            recent_tool,
        ]

        state: Dict[str, Any] = {"messages": messages, "step": 10}

        updated = prune_and_compress_state(state)
        # 应该返回包含更新后 messages 的增量
        assert "messages" in updated

        new_messages: List[Any] = updated["messages"]
        assert len(new_messages) == len(messages)

        # 最近一轮的工具调用不应被 compact（受 "最近两轮 user 对话" 保护）
        new_recent_tool = new_messages[-1]
        assert isinstance(new_recent_tool, ToolMessage)
        assert new_recent_tool.content == recent_tool.content

        # 较早的 old_tool 应该被 compact：content 变成占位文本 + metadata 中有 compacted 标记
        new_old_tool = new_messages[2]
        assert isinstance(new_old_tool, ToolMessage)
        assert new_old_tool.content == "[Old tool result compacted]"
        meta = getattr(new_old_tool, "metadata", {}) or {}
        extra = getattr(new_old_tool, "additional_kwargs", {}) or {}
        assert meta.get("compacted") or extra.get("compacted")

    def test_prune_does_not_compact_protected_tool(self, monkeypatch: pytest.MonkeyPatch):
        """受保护工具（如 skill）不应被 compact。"""

        import langchain_agentx.opencode_style_agent.compaction as comp

        # 降低阈值，确保如果不处理保护逻辑会被 compact，同时避免 Compress 触发
        monkeypatch.setattr(comp, "PRUNE_MINIMUM_TOKENS", 1, raising=False)
        monkeypatch.setattr(comp, "PRUNE_PROTECT_TOKENS", 1, raising=False)
        monkeypatch.setenv("OPENCODE_MAX_CONTEXT_TOKENS", "1000000")
        monkeypatch.setenv("OPENCODE_COMPRESS_TRIGGER_FRACTION", "0.99")

        protected_tool_name = next(iter(PRUNE_PROTECTED_TOOLS))

        tool_msg = _make_tool_message(
            content="z" * 30,
            tool_call_id="skill_call",
            name=protected_tool_name,
        )

        messages: List[Any] = [
            HumanMessage(content="old user"),
            AIMessage(content="old ai"),
            tool_msg,
        ]

        state: Dict[str, Any] = {"messages": messages, "step": 5}

        updated = prune_and_compress_state(state)

        # 因为是受保护工具，应该完全不更新
        assert updated == {}

    def test_prune_protects_skill_when_tool_message_has_no_name(
        self, monkeypatch: pytest.MonkeyPatch
    ):
        """当 ToolMessage 无 name（如 wrap_tool_call 缓存返回）时，通过 tool_call_id 反查 AIMessage 识别 skill。"""

        import langchain_agentx.opencode_style_agent.compaction as comp

        monkeypatch.setattr(comp, "PRUNE_MINIMUM_TOKENS", 1, raising=False)
        monkeypatch.setattr(comp, "PRUNE_PROTECT_TOKENS", 1, raising=False)
        monkeypatch.setenv("OPENCODE_MAX_CONTEXT_TOKENS", "1000000")
        monkeypatch.setenv("OPENCODE_COMPRESS_TRIGGER_FRACTION", "0.99")

        # 模拟无 name 的 ToolMessage（wrap_tool_call 缓存等场景）
        tool_msg = ToolMessage(content="z" * 30, tool_call_id="skill_call")
        assert getattr(tool_msg, "name", None) is None

        # AIMessage 中有 tool_call，用于兜底反查
        ai_msg = AIMessage(
            content="",
            tool_calls=[{"id": "skill_call", "name": "skill", "args": {}}],
        )

        messages: List[Any] = [
            HumanMessage(content="old user"),
            ai_msg,
            tool_msg,
        ]

        state: Dict[str, Any] = {"messages": messages, "step": 5}

        updated = prune_and_compress_state(state)

        # 应通过 tool_call_id 反查识别为 skill，不被 compact
        assert updated == {}

    def test_prune_stops_at_summary_boundary(self, monkeypatch: pytest.MonkeyPatch):
        """遇到 summary 消息应停止向更早历史扫描，不再裁剪 summary 之前的工具输出。"""

        import langchain_agentx.opencode_style_agent.compaction as comp

        # 降低阈值，若不考虑 summary 边界，old_tool 将会被裁剪
        monkeypatch.setattr(comp, "PRUNE_MINIMUM_TOKENS", 1, raising=False)
        monkeypatch.setattr(comp, "PRUNE_PROTECT_TOKENS", 1, raising=False)
        monkeypatch.setenv("OPENCODE_MAX_CONTEXT_TOKENS", "1000000")
        monkeypatch.setenv("OPENCODE_COMPRESS_TRIGGER_FRACTION", "0.99")

        old_tool = _make_tool_message(
            content="x" * 50,
            tool_call_id="old_call",
            name="normal_tool",
        )

        # summary 消息（由 Compress 生成的 HumanMessage），应作为 prune 边界
        summary_msg = HumanMessage(
            content="## Goal\n\nsummary here",
            additional_kwargs={"summary": True},
        )

        recent_tool = _make_tool_message(
            content="y" * 50,
            tool_call_id="recent_call",
            name="normal_tool",
        )

        messages: List[Any] = [
            old_tool,
            summary_msg,
            HumanMessage(content="recent user"),
            recent_tool,
        ]

        state: Dict[str, Any] = {"messages": messages, "step": 3}

        updated = prune_and_compress_state(state)

        # 因为 summary_msg 应该阻止向更早历史扫描，old_tool 不应被裁剪
        # 若逻辑正确，updated 要么为空，要么只可能裁剪 summary 之后的工具（这里也不满足 token 条件）
        if not updated:
            # 完全没有发生裁剪
            assert True
        else:
            new_messages: List[Any] = updated["messages"]
            new_old_tool = new_messages[0]
            assert isinstance(new_old_tool, ToolMessage)
            assert new_old_tool.content == old_tool.content


class TestCompressLogic:
    """测试 Compress 行为（摘要插入与消息截断）。"""

    def test_compress_inserts_summary_and_keeps_recent_messages_mechanical(
        self,
        monkeypatch: pytest.MonkeyPatch,
    ):
        """在不传 model 时，压缩应插入机械结构化 summary，并保留最近若干条消息。"""

        import langchain_agentx.opencode_style_agent.compaction as comp

        # 将 prune 阈值调高，避免本用例被 Prune 干扰（仅测试 Compress）
        monkeypatch.setattr(comp, "PRUNE_MINIMUM_TOKENS", 10_000_000, raising=False)
        monkeypatch.setattr(comp, "PRUNE_PROTECT_TOKENS", 10_000_000, raising=False)

        # 配置极小的上下文阈值，强制触发 Compress
        monkeypatch.setenv("OPENCODE_MAX_CONTEXT_TOKENS", "100")
        monkeypatch.setenv("OPENCODE_COMPRESS_TRIGGER_FRACTION", "0.1")  # 10 tokens
        monkeypatch.setenv("OPENCODE_COMPRESS_KEEP_RECENT_MESSAGES", "2")

        # 构造一段足够长的对话历史（从旧到新）
        h0 = HumanMessage(content="帮我分析代码库架构")  # Goal
        a1 = AIMessage(content="我会先检查项目结构，然后查看主要模块。")
        t1 = _make_tool_message(
            content="read_file('/src/main.py') -> ... 很长的输出 ..." * 5,
            tool_call_id="call-1",
            name="read_file",
        )
        h1 = HumanMessage(content="重点关注认证和权限相关的模块。")
        a2 = AIMessage(content="已找到 auth/ 与 permissions/ 相关模块，正在分析。")
        t2 = _make_tool_message(
            content="grep('auth', path='/src') -> ... 很长的输出 ..." * 5,
            tool_call_id="call-2",
            name="grep",
        )
        # 最近两条消息（应被 keep_recent_messages 保护）
        a_last = AIMessage(content="分析完毕，准备给出总结。")
        t_last = _make_tool_message(
            content="write_report('/tmp/report.md') -> ok",
            tool_call_id="call-3",
            name="write_report",
        )

        messages: List[Any] = [h0, a1, t1, h1, a2, t2, a_last, t_last]
        state: Dict[str, Any] = {"messages": messages, "step": 20}

        updated = prune_and_compress_state(state)

        # Compress 应该生效并返回更新后的 messages
        assert "messages" in updated
        new_messages: List[Any] = updated["messages"]

        # 1) 新消息数应小于原始消息数（前面的历史被折叠为摘要）
        assert len(new_messages) < len(messages)

        # 2) 应存在一个带 summary 标记的 HumanMessage
        summary_msgs = [
            m
            for m in new_messages
            if isinstance(m, HumanMessage)
            and (getattr(m, "additional_kwargs", {}) or {}).get("summary")
        ]
        assert len(summary_msgs) == 1
        summary_msg = summary_msgs[0]
        assert "## Goal" in summary_msg.content  # 机械摘要采用结构化模板

        # 3) 最近 keep_recent_messages 条消息应被完整保留（即最后两条）
        assert isinstance(new_messages[-2], AIMessage)
        assert isinstance(new_messages[-1], ToolMessage)
        assert new_messages[-2].content == a_last.content
        assert new_messages[-1].content == t_last.content

    def test_compress_uses_model_when_provided(
        self,
        monkeypatch: pytest.MonkeyPatch,
    ):
        """当提供 model 时，应优先使用 model.invoke 的内容作为 summary 文本。"""

        import langchain_agentx.opencode_style_agent.compaction as comp

        # 同样屏蔽 Prune，只测试 Compress + model 分支
        monkeypatch.setattr(comp, "PRUNE_MINIMUM_TOKENS", 10_000_000, raising=False)
        monkeypatch.setattr(comp, "PRUNE_PROTECT_TOKENS", 10_000_000, raising=False)

        monkeypatch.setenv("OPENCODE_MAX_CONTEXT_TOKENS", "100")
        monkeypatch.setenv("OPENCODE_COMPRESS_TRIGGER_FRACTION", "0.1")
        monkeypatch.setenv("OPENCODE_COMPRESS_KEEP_RECENT_MESSAGES", "2")

        # 构造与上一个测试类似的消息序列
        h0 = HumanMessage(content="帮我分析代码库架构")
        a1 = AIMessage(content="我会先检查项目结构，然后查看主要模块。")
        t1 = _make_tool_message(
            content="read_file('/src/main.py') -> ... 很长的输出 ..." * 5,
            tool_call_id="call-1",
            name="read_file",
        )
        a_last = AIMessage(content="分析完毕，准备给出总结。")
        t_last = _make_tool_message(
            content="write_report('/tmp/report.md') -> ok",
            tool_call_id="call-2",
            name="write_report",
        )

        messages: List[Any] = [h0, a1, t1, a_last, t_last]
        state: Dict[str, Any] = {"messages": messages, "step": 30}

        # 构造一个 fake model，返回固定的 AIMessage
        class FakeModel:
            def invoke(self, msgs: List[Any]) -> AIMessage:
                return AIMessage(content="## Goal\n\nLLM-based summary\n")

        fake_model = FakeModel()

        updated = prune_and_compress_state(state, model=fake_model)

        assert "messages" in updated
        new_messages: List[Any] = updated["messages"]

        summary_msgs = [
            m
            for m in new_messages
            if isinstance(m, HumanMessage)
            and (getattr(m, "additional_kwargs", {}) or {}).get("summary")
        ]
        assert len(summary_msgs) == 1
        summary_msg = summary_msgs[0]
        # 应该直接采用 fake_model.invoke 返回的内容，而不是机械模板
        assert "LLM-based summary" in summary_msg.content

        # 最后两条消息仍然是原始的 a_last / t_last
        assert isinstance(new_messages[-2], AIMessage)
        assert isinstance(new_messages[-1], ToolMessage)
        assert new_messages[-2].content == a_last.content
        assert new_messages[-1].content == t_last.content

    def test_compress_matches_design_example_structure(
        self,
        monkeypatch: pytest.MonkeyPatch,
    ):
        """按照设计文档中的 H0/skill/plan/tools 队列模拟一次完整 Compress 结构."""

        import langchain_agentx.opencode_style_agent.compaction as comp

        # 只关注 Compress 结构，不让 Prune 介入
        monkeypatch.setattr(comp, "PRUNE_MINIMUM_TOKENS", 10_000_000, raising=False)
        monkeypatch.setattr(comp, "PRUNE_PROTECT_TOKENS", 10_000_000, raising=False)

        # 强制触发 Compress，且只保留最后 2 条消息
        monkeypatch.setenv("OPENCODE_MAX_CONTEXT_TOKENS", "200")
        monkeypatch.setenv("OPENCODE_COMPRESS_TRIGGER_FRACTION", "0.1")
        monkeypatch.setenv("OPENCODE_COMPRESS_KEEP_RECENT_MESSAGES", "2")

        # 按设计文档的顺序构造消息（从旧到新）
        h0 = HumanMessage(content="帮我分析代码库架构")  # H0

        a_skill = AIMessage(
            content="调用 skill: architecture_analysis",
            additional_kwargs={
                "tool_calls": [
                    {"name": "skill", "args": {"skill_name": "architecture_analysis"}}
                ]
            },
        )
        t_skill = _make_tool_message(
            content="SKILL.md full content ...",
            tool_call_id="skill-call",
            name="skill",
        )

        a_plan1 = AIMessage(content="plan v1", additional_kwargs={"tool_calls": [{"name": "write_todos"}]})
        t_plan1 = _make_tool_message(
            content="plan step1, step2 ...",
            tool_call_id="plan-1",
            name="write_todos",
        )

        a_tool1 = AIMessage(content="call tool1", additional_kwargs={"tool_calls": [{"name": "tool1"}]})
        t_tool1 = _make_tool_message(
            content="<tool1 的详细输出>" * 5,
            tool_call_id="tool-1",
            name="tool1",
        )

        a_plan2 = AIMessage(content="plan v2", additional_kwargs={"tool_calls": [{"name": "write_todos"}]})
        t_plan2 = _make_tool_message(
            content="更新后的 plan（含 step2...）",
            tool_call_id="plan-2",
            name="write_todos",
        )

        a_tool2 = AIMessage(content="call tool2", additional_kwargs={"tool_calls": [{"name": "tool2"}]})
        t_tool2 = _make_tool_message(
            content="<tool2 的详细输出>" * 5,
            tool_call_id="tool-2",
            name="tool2",
        )

        # 再追加几轮工具调用，最后一轮为 A_tooln/T_tooln
        a_tooln = AIMessage(content="call toolN", additional_kwargs={"tool_calls": [{"name": "toolN"}]})
        t_tooln = _make_tool_message(
            content="<toolN 的结果>" * 5,
            tool_call_id="tool-N",
            name="toolN",
        )

        messages: List[Any] = [
            h0,
            a_skill,
            t_skill,
            a_plan1,
            t_plan1,
            a_tool1,
            t_tool1,
            a_plan2,
            t_plan2,
            a_tool2,
            t_tool2,
            a_tooln,
            t_tooln,
        ]

        state: Dict[str, Any] = {"messages": messages, "step": 50}

        updated = prune_and_compress_state(state)
        assert "messages" in updated
        new_messages: List[Any] = updated["messages"]

        # 1) 新队列比原始短（中间历史被折叠）
        assert len(new_messages) < len(messages)

        # 2) 结构检查：应包含 H0、T_skill、最新 plan（A_plan2/T_plan2）、summary、以及最后一轮工具（A_tooln/T_tooln）
        #    A_skill 在文档中是“可选保留”，当前实现不强制保留，因此这里不做断言。
        # from old to new:
        #   [H0, T_skill, A_plan2, T_plan2, H_summary, A_tooln, T_tooln]

        # H0
        assert isinstance(new_messages[0], HumanMessage)
        assert new_messages[0].content == h0.content

        # T_skill 应该保留在早期位置
        assert any(isinstance(m, ToolMessage) and getattr(m, "name", None) == "skill" for m in new_messages)

        # 最新 plan（A_plan2/T_plan2）应仍然保留
        assert any(isinstance(m, AIMessage) and m.content == a_plan2.content for m in new_messages)
        assert any(
            isinstance(m, ToolMessage)
            and getattr(m, "name", None) == "write_todos"
            and m.content == t_plan2.content
            for m in new_messages
        )

        # summary 消息
        summary_msgs = [
            m
            for m in new_messages
            if isinstance(m, HumanMessage)
            and (getattr(m, "additional_kwargs", {}) or {}).get("summary")
        ]
        assert len(summary_msgs) == 1

        # 最后两条应为 A_tooln/T_tooln
        assert isinstance(new_messages[-2], AIMessage)
        assert isinstance(new_messages[-1], ToolMessage)
        assert new_messages[-2].content == a_tooln.content
        assert new_messages[-1].content == t_tooln.content


if __name__ == "__main__":
    pytest.main([__file__, "-q"])

