"""
测试 message_utils 工具模块

覆盖：
- estimate_tokens: 含 tool_calls / tool_call_id / reasoning_content
- split_into_rounds: round 边界识别，空列表等边界情况
- build_compact_view: 不足预算 / 超出预算 / round 完整性
- validate_tool_call_integrity: 合法 / orphan tool_calls / orphan tool_messages
"""
import sys
from pathlib import Path

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

import pytest
from aacode.utils.message_utils import (
    estimate_tokens,
    split_into_rounds,
    build_compact_view,
    validate_tool_call_integrity,
)


# ---------------------------------------------------------------
# estimate_tokens
# ---------------------------------------------------------------

class TestEstimateTokens:
    def test_empty_list(self):
        assert estimate_tokens(None, []) == 0

    def test_content_only_no_encoding(self):
        msgs = [{"role": "user", "content": "0123456789"}]
        # 10 chars // 4 = 2
        assert estimate_tokens(None, msgs) == 2

    def test_content_only_no_encoding_empty_content(self):
        msgs = [{"role": "user", "content": ""}]
        assert estimate_tokens(None, msgs) == 0

    def test_counts_tool_calls(self):
        msgs = [
            {
                "role": "assistant",
                "content": "a",
                "tool_calls": [
                    {
                        "id": "call_1",
                        "type": "function",
                        "function": {
                            "name": "run_shell",
                            "arguments": '{"path": "foo.py"}',
                        },
                    }
                ],
            }
        ]
        tokens_with = estimate_tokens(None, msgs)
        # content only would be len("a")//4 = 0
        # with tool_calls JSON should be > 0
        assert tokens_with > 0
        # verify tool_calls adds extra tokens beyond content
        content_only = estimate_tokens(None, [{"role": "user", "content": "a"}])
        assert tokens_with > content_only

    def test_counts_tool_call_id(self):
        msgs = [{"role": "tool", "tool_call_id": "call_abc_123", "content": "ok"}]
        tokens_with = estimate_tokens(None, msgs)
        content_only = estimate_tokens(None, [{"role": "tool", "content": "ok"}])
        # "call_abc_123" should add tokens
        assert tokens_with > content_only

    def test_counts_reasoning_content(self):
        msgs = [
            {
                "role": "assistant",
                "content": "x",
                "reasoning_content": "Let me think about this carefully...",
            }
        ]
        tokens_with = estimate_tokens(None, msgs)
        content_only = estimate_tokens(None, [{"role": "assistant", "content": "x"}])
        assert tokens_with > content_only

    def test_multiple_fields_summed(self):
        msgs = [
            {
                "role": "assistant",
                "content": "Hello world",
                "tool_calls": [
                    {
                        "id": "c1",
                        "type": "function",
                        "function": {"name": "f", "arguments": "{}"},
                    }
                ],
                "reasoning_content": "thinking...",
            },
            {"role": "tool", "tool_call_id": "c1", "content": "result"},
        ]
        tokens = estimate_tokens(None, msgs)
        assert tokens > 5  # reasonably > content-only estimate


# ---------------------------------------------------------------
# split_into_rounds
# ---------------------------------------------------------------

def _asst(text="a", tool_calls=None):
    return {"role": "assistant", "content": text, "tool_calls": tool_calls or []}


def _tool(call_id, content="r"):
    return {"role": "tool", "tool_call_id": call_id, "content": content}


def _user(text="q"):
    return {"role": "user", "content": text}


def _sys(text="s"):
    return {"role": "system", "content": text}


class TestSplitIntoRounds:
    def test_empty(self):
        assert split_into_rounds([]) == []

    def test_single_message(self):
        rounds = split_into_rounds([_user("hello")])
        assert len(rounds) == 1
        assert len(rounds[0]) == 1

    def test_system_and_user_separate_rounds(self):
        rounds = split_into_rounds([_sys(), _user()])
        assert len(rounds) == 2

    def test_assistant_with_tools_one_round(self):
        msgs = [
            _asst("a1", [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": "{}"}}]),
            _tool("c1"),
            _tool("c2"),  # same round even if id doesn't match (controlled by role)
        ]
        rounds = split_into_rounds(msgs)
        assert len(rounds) == 1, f"Expected 1 round, got {len(rounds)}"
        assert len(rounds[0]) == 3

    def test_multiple_tool_rounds(self):
        msgs = [
            _asst("a1", [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": "{}"}}]),
            _tool("c1"),
            _asst("a2", [{"id": "c2", "type": "function", "function": {"name": "g", "arguments": "{}"}}]),
            _tool("c2"),
        ]
        rounds = split_into_rounds(msgs)
        assert len(rounds) == 2
        assert len(rounds[0]) == 2
        assert len(rounds[1]) == 2

    def test_user_then_assistant_then_tool(self):
        msgs = [_user("q"), _asst("a", [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": "{}"}}]), _tool("c1")]
        rounds = split_into_rounds(msgs)
        assert len(rounds) == 2  # user, asst+tool
        assert len(rounds[0]) == 1
        assert len(rounds[1]) == 2

    def test_round_trip_preserves(self):
        msgs = [
            _sys(),
            _user("task"),
            _asst("a1", [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": "{}"}}]),
            _tool("c1"),
            _asst("a2"),
            _user("next"),
            _asst("a3", [{"id": "c3", "type": "function", "function": {"name": "h", "arguments": "{}"}}]),
            _tool("c3"),
        ]
        rounds = split_into_rounds(msgs)
        flat = [m for r in rounds for m in r]
        assert flat == msgs

    def test_consecutive_user_messages(self):
        """连续 user 消息各自独立成 round"""
        msgs = [_user("q1"), _user("q2")]
        rounds = split_into_rounds(msgs)
        assert len(rounds) == 2

    def test_tool_only_starts_own_round(self):
        """如果消息以 tool 开头，它自成一轮"""
        msgs = [_tool("cid"), _user("q")]
        rounds = split_into_rounds(msgs)
        assert len(rounds) == 2
        assert len(rounds[0]) == 1
        assert rounds[0][0]["role"] == "tool"


# ---------------------------------------------------------------
# build_compact_view
# ---------------------------------------------------------------

class TestBuildCompactView:
    def test_under_budget_returns_full(self):
        msgs = [_sys(), _user("task"), _asst(), _tool("c1")]
        view, compacted, _ = build_compact_view(None, msgs, 1000000)
        assert compacted is False
        assert len(view) == len(msgs)
        assert view == msgs  # returns a copy but equal

    def test_exact_split_works(self):
        """当 rounds 数恰好等于 protect+keep 时无需压缩"""
        msgs = [_sys(), _user(), _asst(), _tool("c1"), _asst(), _tool("c2")]
        # 4 rounds: [sys], [user], [asst+tool], [asst+tool]
        view, compacted, _ = build_compact_view(
            None, msgs, max_tokens=1, protect_first_rounds=2, keep_last_rounds=2
        )
        assert compacted is False  # exactly 4 rounds, no need to compact

    def test_over_budget_compacts(self):
        msgs = [
            _sys(),
            _asst("a0", [{"id": "c0", "type": "function", "function": {"name": "init", "arguments": "{}"}}]),
            _tool("c0"),
            _asst("a1", [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": "{}"}}]),
            _tool("c1"),
            _user("task"),
            _asst("a2", [{"id": "c2", "type": "function", "function": {"name": "g", "arguments": "{}"}}]),
            _tool("c2"),
        ]
        # rounds: [sys], [a0+t0], [a1+t1], [user], [a2+t2] = 5 rounds
        # user at idx 3, pre_user=[sys, a0+t0, a1+t1], protect_first=2
        # keep 2 of 3, skip 1 → compacted=True
        view, compacted, _ = build_compact_view(
            None, msgs, max_tokens=1, protect_first_rounds=2, keep_last_rounds=1
        )
        assert compacted is True
        assert len(view) < len(msgs)

    def test_compacted_view_passes_integrity(self):
        """压缩视图中的 tool_calls/tool_messages 完整配对"""
        msgs = [
            _sys(),
            _asst("a0", [
                {"id": "c0a", "type": "function", "function": {"name": "f", "arguments": "{}"}},
                {"id": "c0b", "type": "function", "function": {"name": "g", "arguments": "{}"}},
            ]),
            _tool("c0a"),
            _tool("c0b"),
            _user("task"),
            _asst("a1", [{"id": "c1", "type": "function", "function": {"name": "h", "arguments": "{}"}}]),
            _tool("c1"),
        ]
        # rounds: [sys], [a0+t0a+t0b], [user], [a1+t1] = 4 rounds
        # user at idx 2, pre_user=[sys, a0+t0a+t0b], protect_first=1
        # keep 1 of 2 → skipped=1 → compacted=True
        view, compacted, _ = build_compact_view(
            None, msgs, max_tokens=1, protect_first_rounds=1, keep_last_rounds=1
        )
        assert compacted is True
        # should not raise
        validate_tool_call_integrity(view)

    def test_with_cached_summary(self):
        msgs = [
            _sys(),
            _asst("a0", [{"id": "c0", "type": "function", "function": {"name": "f", "arguments": "{}"}}]),
            _tool("c0"),
            _user("task"),
            _asst("a1", [{"id": "c1", "type": "function", "function": {"name": "g", "arguments": "{}"}}]),
            _tool("c1"),
        ]
        # rounds: [sys], [a0+t0], [user], [a1+t1] = 4 rounds
        # user at idx 2, pre_user=[sys, a0+t0], protect_first=1
        # keep 1 of 2 → skipped=1 → compacted=True
        view, compacted, _ = build_compact_view(
            None, msgs, max_tokens=1,
            protect_first_rounds=1, keep_last_rounds=1,
            cached_summary="This is my custom summary of earlier rounds",
        )
        assert compacted is True
        # ensure cached summary is present
        summaries = [m for m in view if "custom summary" in str(m.get("content", ""))]
        assert len(summaries) == 1

    def test_original_messages_not_modified(self):
        msgs = [_sys(), _user(), _asst(), _tool("c1")]
        original_len = len(msgs)
        original_first = msgs[0]
        view, _, _ = build_compact_view(None, msgs, max_tokens=1, protect_first_rounds=1, keep_last_rounds=1)
        # original list should be unchanged
        assert len(msgs) == original_len
        assert msgs[0] is original_first
        # view should be a different list
        assert view is not msgs

    def test_returns_copy_when_under_budget(self):
        msgs = [_user("q")]
        view, _, _ = build_compact_view(None, msgs, 1000000)
        assert view == msgs
        assert view is not msgs  # returns a copy


# ---------------------------------------------------------------
# validate_tool_call_integrity
# ---------------------------------------------------------------

class TestValidateToolCallIntegrity:
    def test_empty_ok(self):
        validate_tool_call_integrity([])

    def test_no_tool_calls_ok(self):
        msgs = [_user("q"), _asst("a")]
        validate_tool_call_integrity(msgs)

    def test_valid_pairing_ok(self):
        msgs = [
            _asst("a", [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": "{}"}}]),
            _tool("c1"),
        ]
        validate_tool_call_integrity(msgs)

    def test_multiple_tool_calls_one_assistant(self):
        msgs = [
            _asst("a", [
                {"id": "c1", "type": "function", "function": {"name": "f", "arguments": "{}"}},
                {"id": "c2", "type": "function", "function": {"name": "g", "arguments": "{}"}},
            ]),
            _tool("c1"),
            _tool("c2"),
        ]
        validate_tool_call_integrity(msgs)

    def test_orphaned_tool_calls_raises(self):
        msgs = [
            _asst("a", [{"id": "orphan", "type": "function", "function": {"name": "f", "arguments": "{}"}}]),
        ]
        with pytest.raises(ValueError, match="Orphaned tool_calls"):
            validate_tool_call_integrity(msgs)

    def test_orphaned_tool_message_raises(self):
        msgs = [_tool("no_such_call_id")]
        with pytest.raises(ValueError, match="Orphaned tool message"):
            validate_tool_call_integrity(msgs)

    def test_partial_orphan_in_mixed(self):
        """一个 tool_call 有对应 tool，另一个没有"""
        msgs = [
            _asst("a", [
                {"id": "c1", "type": "function", "function": {"name": "f", "arguments": "{}"}},
                {"id": "c2", "type": "function", "function": {"name": "g", "arguments": "{}"}},
            ]),
            _tool("c1"),
            # c2 missing
        ]
        with pytest.raises(ValueError, match="Orphaned tool_calls"):
            validate_tool_call_integrity(msgs)

    def test_interleaved_pairs_ok(self):
        msgs = [
            _asst("a1", [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": "{}"}}]),
            _tool("c1"),
            _asst("a2", [{"id": "c2", "type": "function", "function": {"name": "g", "arguments": "{}"}}]),
            _tool("c2"),
        ]
        validate_tool_call_integrity(msgs)

    def test_unexpected_tool_before_any_assistant(self):
        msgs = [_tool("c99")]
        with pytest.raises(ValueError, match="Orphaned tool message"):
            validate_tool_call_integrity(msgs)
