"""
专门测试：压缩触发逻辑 —— 确保以压缩后实际送入模型的 token 数为准

验证要点：
1. 触发检查基于 compact_view_tokens（压缩视图），而非全量消息 total_tokens
2. 压缩后 _post_compact_tokens 正确记录基准值
3. 压缩后 msg_count_changed < 20 时不会再次触发压缩（防次次触发）
4. 累积 20+ 条新消息后可再次触发压缩
5. 模型调用（_build_compact_view 返回的 view）与触发检查使用同一 token 数
6. protect_last_user_rounds=2 在完整压缩流程中正确保护最后 2 轮用户消息
"""
import asyncio
import sys
import tempfile
from pathlib import Path
from unittest.mock import AsyncMock, patch

sys.path.insert(0, str(Path(__file__).parent.parent))

import pytest
from aacode.core.react_loop import AsyncReActLoop
from aacode.utils.context_manager import ContextManager
from aacode.utils.message_utils import (
    estimate_tokens,
    validate_tool_call_integrity,
    build_compact_view,
    split_into_rounds,
    _find_last_n_real_user_round,
)


# ── helpers ──────────────────────────────────────────────────────

def _make_loop(**kwargs):
    """创建 AsyncReActLoop 实例"""
    td = tempfile.mkdtemp()
    from aacode.config import ContextConfig

    cfg = ContextConfig()
    for k, v in kwargs.items():
        if hasattr(cfg, k):
            setattr(cfg, k, v)

    async def dummy_caller(messages):
        return {"text": "iteration complete", "tool_calls": []}

    cm = ContextManager(Path(td))
    loop = AsyncReActLoop(
        model_caller=dummy_caller,
        tools={},
        context_manager=cm,
        context_config=cfg,
    )
    return loop, td


def _build_messages(iterations=10, double_size=False):
    """构建带 tool_calls 和用户消息的多轮消息，确保有可压缩的 pre-user 内容"""
    size = 500 if not double_size else 2000
    msgs = [{"role": "system", "content": "sys prompt " + "x" * 100}]

    # 用户消息之前的轮次（可被压缩）
    for i in range(5):
        tc_id = f"pre_{i}"
        msgs.append({
            "role": "assistant",
            "content": f"pre_a{i} " + "p" * size,
            "tool_calls": [{
                "id": tc_id, "type": "function",
                "function": {"name": "run_shell", "arguments": f'{{"cmd":"pre {i}"}}'},
            }],
        })
        msgs.append({
            "role": "tool",
            "tool_call_id": tc_id,
            "content": f"pre_r{i} " + "q" * size,
        })

    # 用户消息（保护边界起点）
    msgs.append({"role": "user", "content": "Main task: do the work " + "t" * size})

    # 用户消息之后的轮次（不应被压缩）
    for i in range(iterations):
        tc_id = f"c{i}"
        msgs.append({
            "role": "assistant",
            "content": f"thinking {i} " + "a" * size,
            "tool_calls": [{
                "id": tc_id, "type": "function",
                "function": {"name": "run_shell", "arguments": f'{{"cmd":"echo {i}"}}'},
            }],
        })
        msgs.append({
            "role": "tool",
            "tool_call_id": tc_id,
            "content": f"result {i} " + "r" * size,
        })
    return msgs


def _build_multi_user_messages():
    """构建多用户消息：3 个 user round，每个后面跟若干 assistant+tool 轮"""
    msgs = [{"role": "system", "content": "sys"}]
    # user 1 + 3 轮 asst+tool
    msgs.append({"role": "user", "content": "Task 1: first task"})
    for i in range(3):
        tc_id = f"c1_{i}"
        msgs.append({
            "role": "assistant", "content": f"a1_{i} " + "z" * 800,
            "tool_calls": [{
                "id": tc_id, "type": "function",
                "function": {"name": "f", "arguments": "{}"},
            }],
        })
        msgs.append({"role": "tool", "tool_call_id": tc_id, "content": f"r1_{i} " + "z" * 800})
    # user 2 + 3 轮 asst+tool
    msgs.append({"role": "user", "content": "Task 2: second task"})
    for i in range(3):
        tc_id = f"c2_{i}"
        msgs.append({
            "role": "assistant", "content": f"a2_{i} " + "z" * 800,
            "tool_calls": [{
                "id": tc_id, "type": "function",
                "function": {"name": "f", "arguments": "{}"},
            }],
        })
        msgs.append({"role": "tool", "tool_call_id": tc_id, "content": f"r2_{i} " + "z" * 800})
    # user 3 + 1 轮 asst+tool
    msgs.append({"role": "user", "content": "Task 3: latest task"})
    tc_id = "c3_0"
    msgs.append({
        "role": "assistant", "content": f"a3_0 " + "z" * 800,
        "tool_calls": [{
            "id": tc_id, "type": "function",
            "function": {"name": "f", "arguments": "{}"},
        }],
    })
    msgs.append({"role": "tool", "tool_call_id": tc_id, "content": f"r3_0 " + "z" * 800})
    return msgs


# ═══════════════════════════════════════════════════════════════
# 1. 触发检查基于 compact view token 数
# ═══════════════════════════════════════════════════════════════

class TestTriggerBasedOnCompactView:
    """验证触发检查使用的是压缩视图 token 数，而非全量消息总 token"""

    def test_build_compact_view_returns_compact_tokens(self):
        """_build_compact_view 返回的是压缩后的 token 数"""
        loop, _ = _make_loop(
            max_context_length=1000,
            compact_trigger_tokens=800,
        )
        msgs = _build_messages(iterations=10)
        full_tokens = estimate_tokens(loop.encoding, msgs)

        view, was_compacted, compact_tokens = loop._build_compact_view(msgs)

        # 压缩视图 token 数应小于全量 token 数
        assert compact_tokens < full_tokens, (
            f"Compact view tokens ({compact_tokens}) should be less than full tokens ({full_tokens})"
        )
        # 压缩视图应与传入模型的消息一致
        assert view is not msgs, "Should return new list, not original"

    def test_trigger_check_uses_compact_view_tokens(self):
        """每次迭代的触发检查调用 _build_compact_view，返回 compact 后的 token"""
        loop, _ = _make_loop(
            max_context_length=2000,
            compact_trigger_tokens=1500,
        )
        msgs = _build_messages(iterations=10)

        # 全量 token 应远超触发阈值
        full_tokens = estimate_tokens(loop.encoding, msgs)
        assert full_tokens > 1500, f"Full tokens {full_tokens} should exceed trigger 1500"

        # 但 _build_compact_view 返回的 compact_view_tokens 可能低于阈值
        view, was_compacted, compact_view_tokens = loop._build_compact_view(msgs)
        # 验证 compact_view_tokens 确实来自压缩视图
        view_tokens = estimate_tokens(loop.encoding, view)
        assert compact_view_tokens == view_tokens, (
            f"compact_view_tokens ({compact_view_tokens}) should equal actual view tokens ({view_tokens})"
        )


# ═══════════════════════════════════════════════════════════════
# 2. _post_compact_tokens 基准值
# ═══════════════════════════════════════════════════════════════

class TestPostCompactTokens:
    """验证压缩后 _post_compact_tokens 正确记录"""

    def test_post_compact_tokens_set_after_compaction(self):
        """模拟 _compact_context 后 _post_compact_tokens 被正确设置"""
        loop, _ = _make_loop(
            max_context_length=2000,
            compact_trigger_tokens=1500,
            compact_keep_rounds=2,
        )
        msgs = _build_messages(iterations=10)

        # 初始为 0
        assert loop._post_compact_tokens == 0

        # 运行压缩
        loop._compact_summary_msg_count = 0
        asyncio.run(loop._compact_context(msgs))

        # 压缩后应设置非零值
        assert loop._post_compact_tokens > 0, (
            f"_post_compact_tokens should be > 0 after compaction, got {loop._post_compact_tokens}"
        )

    def test_post_compact_tokens_less_than_full(self):
        """压缩后的 token 基准值应小于全量 token"""
        loop, _ = _make_loop(
            max_context_length=2000,
            compact_trigger_tokens=1500,
        )
        msgs = _build_messages(iterations=10)
        full_tokens = estimate_tokens(loop.encoding, msgs)

        loop._compact_summary_msg_count = 0
        asyncio.run(loop._compact_context(msgs))

        assert loop._post_compact_tokens < full_tokens, (
            f"_post_compact_tokens ({loop._post_compact_tokens}) should be less than full ({full_tokens})"
        )

    def test_post_compact_tokens_close_to_build_compact_view(self):
        """_post_compact_tokens 应与 _build_compact_view 直接返回的 token 数一致"""
        loop, _ = _make_loop(
            max_context_length=2000,
            compact_trigger_tokens=1500,
        )
        msgs = _build_messages(iterations=10)

        loop._compact_summary_msg_count = 0
        asyncio.run(loop._compact_context(msgs))

        _, _, tokens_from_build = loop._build_compact_view(msgs)
        assert loop._post_compact_tokens == tokens_from_build, (
            f"_post_compact_tokens ({loop._post_compact_tokens}) should match "
            f"_build_compact_view result ({tokens_from_build})"
        )


# ═══════════════════════════════════════════════════════════════
# 3. 防重复触发（msg_count_changed guard）
# ═══════════════════════════════════════════════════════════════

class TestGuardAgainstReTrigger:
    """验证压缩后不会每次迭代都触发"""

    def test_guard_suppresses_retrigger_within_20_msgs(self):
        """消息增长 < 20 条时，即使超阈值也不重新压缩"""
        loop, _ = _make_loop(
            max_context_length=2000,
            compact_trigger_tokens=1500,
            compact_keep_rounds=2,
        )
        msgs = _build_messages(iterations=10)

        # 首次压缩
        loop._compact_summary = None
        loop._compact_summary_msg_count = 0
        asyncio.run(loop._compact_context(msgs))
        assert loop._compact_summary is not None

        compact_count = 1

        # 模拟追加少量消息后检查触发逻辑
        for _ in range(3):
            msgs.append({"role": "assistant", "content": "extra"})
            msgs.append({"role": "tool", "tool_call_id": "extra", "content": "extra"})

            # msg_count_changed = 新增条数，应 < 20
            msg_count_changed = len(msgs) - loop._compact_summary_msg_count
            assert msg_count_changed < 20, f"msg_count_changed {msg_count_changed} should be < 20"

            # summary 已存在，msg_count_changed < 20 → 不应触发压缩
            if loop._compact_summary is not None and msg_count_changed <= 20:
                pass  # guard 生效，不触发
            else:
                compact_count += 1

        # 应只压缩了 1 次
        assert compact_count == 1, f"Expected 1 compaction, got {compact_count}"

    def test_guard_allows_retrigger_after_20_msgs(self):
        """消息增长 ≥ 20 条时可重新触发压缩"""
        loop, _ = _make_loop(
            max_context_length=5000,
            compact_trigger_tokens=3000,
            compact_keep_rounds=1,
            compact_protect_first_rounds=1,
        )
        msgs = _build_messages(iterations=5, double_size=True)

        # 首次压缩
        loop._compact_summary = None
        loop._compact_summary_msg_count = 0
        asyncio.run(loop._compact_context(msgs))
        assert loop._compact_summary is not None, "First compaction should set summary"

        # 追加 22 条新消息（11 轮 asst+tool）× 大尺寸
        for i in range(11):
            tc_id = f"new_{i}"
            msgs.append({
                "role": "assistant", "content": f"new_a{i} " + "z" * 2000,
                "tool_calls": [{
                    "id": tc_id, "type": "function",
                    "function": {"name": "f", "arguments": "{}"},
                }],
            })
            msgs.append({"role": "tool", "tool_call_id": tc_id, "content": f"new_r{i} " + "z" * 2000})

        msg_count_changed = len(msgs) - loop._compact_summary_msg_count
        assert msg_count_changed > 20, (
            f"Expected > 20 new messages, got {msg_count_changed}"
        )

        # 此时应允许再次压缩
        _, _, compact_view_tokens = loop._build_compact_view(msgs)
        trigger = min(loop.context_config.compact_trigger_tokens, loop.context_config.max_context_length)

        assert compact_view_tokens > trigger, (
            f"compact_view_tokens ({compact_view_tokens}) should exceed trigger ({trigger})"
        )
        # guard 允许通过（msg_count_changed > 20）
        assert msg_count_changed > 20


# ═══════════════════════════════════════════════════════════════
# 4. 模型调用与触发检查 token 一致性
# ═══════════════════════════════════════════════════════════════

class TestTokenConsistency:
    """验证模型调用使用的 token 数与触发检查一致"""

    def test_cached_compact_tokens_matches_build_result(self):
        """_cached_compact_tokens 与 _build_compact_view 返回一致"""
        loop, _ = _make_loop(
            max_context_length=2000,
            compact_trigger_tokens=1500,
        )
        msgs = _build_messages(iterations=10)

        # 模拟模型调用前的缓存（与 react_loop 迭代中 line 235-238 相同逻辑）
        effective_msgs, _, compact_tokens = loop._build_compact_view(msgs)
        loop._cached_compact_view = effective_msgs
        loop._cached_compact_tokens = compact_tokens
        loop._cached_compact_msg_count = len(msgs)

        # 缓存的 token 数应等于视图实际 token 数
        view_tokens = estimate_tokens(loop.encoding, effective_msgs)
        assert loop._cached_compact_tokens == view_tokens, (
            f"Cached tokens ({loop._cached_compact_tokens}) should match view tokens ({view_tokens})"
        )

    def test_trigger_and_model_call_use_same_function(self):
        """触发检查和模型调用都通过 _build_compact_view 计算 token"""
        loop, _ = _make_loop(
            max_context_length=2000,
            compact_trigger_tokens=1500,
        )
        msgs = _build_messages(iterations=5)

        # 模型调用时（line 235）
        _, _, model_tokens = loop._build_compact_view(msgs)

        # 追加一轮完整的 asst(tool_calls) + tool 后触发检查（line 560）
        tc_id = "extra_test"
        msgs.append({
            "role": "assistant", "content": "extra",
            "tool_calls": [{
                "id": tc_id, "type": "function",
                "function": {"name": "f", "arguments": "{}"},
            }],
        })
        msgs.append({"role": "tool", "tool_call_id": tc_id, "content": "extra result"})
        _, _, trigger_tokens = loop._build_compact_view(msgs)

        # 触发检查的 token 数应 ≥ 模型调用时的（因为多了一轮消息）
        assert trigger_tokens >= model_tokens, (
            f"trigger_tokens ({trigger_tokens}) should be >= model_tokens ({model_tokens})"
        )


# ═══════════════════════════════════════════════════════════════
# 5. protect_last_user_rounds=2 集成测试
# ═══════════════════════════════════════════════════════════════

class TestProtectLast2UserRoundsIntegration:
    """验证 protect_last_user_rounds=2 在完整压缩流程中正确"""

    def test_last_2_users_protected_after_compaction(self):
        """压缩后最后 2 条用户消息及其附带消息完整保留"""
        loop, _ = _make_loop(
            max_context_length=2000,
            compact_trigger_tokens=1500,
            compact_keep_rounds=2,
            compact_protect_first_rounds=1,
            compact_protect_user_rounds=2,
        )
        msgs = _build_multi_user_messages()

        # 运行压缩
        loop._compact_summary_msg_count = 0
        asyncio.run(loop._compact_context(msgs))

        # 验证压缩视图
        view, compacted, _ = loop._build_compact_view(msgs)
        assert compacted is True
        validate_tool_call_integrity(view)

        user_texts = [m["content"] for m in view if m["role"] == "user"]
        assert any("Task 2" in t for t in user_texts), "Task 2 (2nd-to-last user) should be preserved"
        assert any("Task 3" in t for t in user_texts), "Task 3 (latest user) should be preserved"

        # Task 1（最早的 user）可能被压缩
        task1_in_view = any("Task 1" in t for t in user_texts)
        if task1_in_view:
            # 如果 Task 1 也在视图中，说明 pre_user 太小没被压缩
            rounds = split_into_rounds(msgs)
            idx = _find_last_n_real_user_round(rounds, 2)
            pre_user = rounds[:idx]
            assert len(pre_user) <= 1, f"pre_user too large: {len(pre_user)} rounds"

    def test_compact_context_uses_n2_split(self):
        """_compact_context 使用 _find_last_n_real_user_round(rounds, 2) 确定保护边界"""
        loop, _ = _make_loop(
            max_context_length=2000,
            compact_trigger_tokens=1500,
            compact_protect_first_rounds=1,
            compact_protect_user_rounds=2,
        )
        msgs = _build_multi_user_messages()

        rounds = split_into_rounds(msgs)
        # 验证保护边界：最后 2 条用户消息的较早者
        last_2_idx = _find_last_n_real_user_round(rounds, 2)
        # 应该有 3 条用户消息：Task 1, Task 2, Task 3
        last_1_idx = _find_last_n_real_user_round(rounds, 1)
        assert last_2_idx < last_1_idx, (
            f"last_2_idx ({last_2_idx}) should be before last_1_idx ({last_1_idx})"
        )

        # 保护边界之后应包含 Task 2 和 Task 3
        post_user = rounds[last_2_idx:]
        post_user_texts = []
        for r in post_user:
            for m in r:
                if m["role"] == "user":
                    post_user_texts.append(m["content"])
        assert any("Task 3" in t for t in post_user_texts), (
            f"Task 3 should be in protected rounds, got {post_user_texts}"
        )
        assert any("Task 2" in t for t in post_user_texts), (
            f"Task 2 should be in protected rounds, got {post_user_texts}"
        )

    def test_tool_integrity_preserved_with_n2(self):
        """protect_last_user_rounds=2 不破坏 tool_calls 完整性"""
        loop, _ = _make_loop(
            max_context_length=2000,
            compact_trigger_tokens=1500,
            compact_protect_first_rounds=1,
            compact_protect_user_rounds=2,
        )
        msgs = _build_multi_user_messages()

        # 运行压缩并验证视图完整性
        loop._compact_summary_msg_count = 0
        asyncio.run(loop._compact_context(msgs))
        view, _, _ = loop._build_compact_view(msgs)
        validate_tool_call_integrity(view)

        # 再次加一轮消息仍应完整
        tc_id = "extra_test"
        msgs.append({
            "role": "assistant", "content": "extra",
            "tool_calls": [{
                "id": tc_id, "type": "function",
                "function": {"name": "f", "arguments": "{}"},
            }],
        })
        msgs.append({"role": "tool", "tool_call_id": tc_id, "content": "extra result"})

        view2, _, _ = loop._build_compact_view(msgs)
        validate_tool_call_integrity(view2)


# ═══════════════════════════════════════════════════════════════
# 6. 边界情况
# ═══════════════════════════════════════════════════════════════

class TestEdgeCases:
    """边界和极端情况"""

    def test_trigger_not_exceeded_no_compaction(self):
        """token 未超阈值时不触发压缩"""
        loop, _ = _make_loop(
            max_context_length=100000,
            compact_trigger_tokens=100000,
        )
        msgs = _build_messages(iterations=3)

        _, was_compacted, compact_view_tokens = loop._build_compact_view(msgs)

        # full tokens <= max_tokens，不压缩
        assert was_compacted is False
        assert compact_view_tokens < 100000

    def test_initial_compact_summary_none(self):
        """初始状态 _compact_summary 为 None"""
        loop, _ = _make_loop()
        assert loop._compact_summary is None
        assert loop._compact_summary_msg_count == 0
        assert loop._post_compact_tokens == 0

    def test_summary_persists_after_compaction(self):
        """压缩后 summary 不为 None"""
        loop, _ = _make_loop(
            max_context_length=2000,
            compact_trigger_tokens=1500,
        )
        msgs = _build_messages(iterations=10)

        assert loop._compact_summary is None
        loop._compact_summary_msg_count = 0
        asyncio.run(loop._compact_context(msgs))
        assert loop._compact_summary is not None
        assert "Smart History Summary" in loop._compact_summary

    def test_msg_count_matches_delta(self):
        """_compact_summary_msg_count 与消息计数差正确"""
        loop, _ = _make_loop(
            max_context_length=2000,
            compact_trigger_tokens=1500,
        )
        msgs = _build_messages(iterations=10)
        original_len = len(msgs)

        loop._compact_summary_msg_count = 0
        asyncio.run(loop._compact_context(msgs))

        assert loop._compact_summary_msg_count == original_len, (
            f"_compact_summary_msg_count ({loop._compact_summary_msg_count}) should equal {original_len}"
        )
