import pytest

from context.base import ContextConfig
from context.default import DefaultContext
from provider import Message, MessageRole, ToolCall, ToolResult


class FakeChunk:
    def __init__(self, text: str):
        self.text = text


class FakeChat:
    def __init__(self, chunks):
        self.chunks = chunks
        self.calls = []

    async def __call__(self, messages, system_prompts, tools=None):
        self.calls.append(
            {"messages": messages, "system_prompts": system_prompts, "tools": tools}
        )
        for text in self.chunks:
            yield FakeChunk(text)


async def empty_chat_fn(*_args, **_kwargs):
    if False:
        yield None


def _tool_result_msg(content: str, tool_content: str, tool_id: str) -> Message:
    return Message(
        role=MessageRole.ASSISTANT,
        content=content,
        tool_result=ToolResult(
            tool_call_id=tool_id,
            name="dummy_tool",
            content=tool_content,
            is_error=False,
        ),
    )


def _tool_call_msg(content: str, tool_id: str) -> Message:
    return Message(
        role=MessageRole.ASSISTANT,
        content=content,
        tool_calls=[ToolCall(id=tool_id, name="dummy_tool", arguments={"x": 1})],
    )


def test_prune_tool_outputs_replaces_old_outputs():
    config = ContextConfig(
        max_tokens=100,
        prune_keep_recent=0,
        prune_protect_chars=0,
        prune_min_chars=1,
    )
    ctx = DefaultContext(empty_chat_fn, config=config)

    messages = [
        Message(role=MessageRole.USER, content="u1"),
        _tool_result_msg("a1", "old-output", "t1"),
        Message(role=MessageRole.USER, content="u2"),
        _tool_result_msg("a2", "mid-output", "t2"),
        Message(role=MessageRole.USER, content="u3"),
        _tool_result_msg("a3", "new-output", "t3"),
    ]

    pruned = ctx._prune_tool_outputs(messages)

    assert pruned[1].tool_result.content == ctx.TOOL_OUTPUT_PLACEHOLDER
    assert pruned[3].tool_result.content == "mid-output"
    assert pruned[5].tool_result.content == "new-output"


def test_prune_tool_outputs_skips_when_below_min_chars():
    config = ContextConfig(
        max_tokens=100,
        prune_keep_recent=0,
        prune_protect_chars=0,
        prune_min_chars=999,
    )
    ctx = DefaultContext(empty_chat_fn, config=config)

    messages = [
        Message(role=MessageRole.USER, content="u1"),
        _tool_result_msg("a1", "small-output", "t1"),
        Message(role=MessageRole.USER, content="u2"),
        _tool_result_msg("a2", "keep-output", "t2"),
        Message(role=MessageRole.USER, content="u3"),
        _tool_result_msg("a3", "keep-output-2", "t3"),
    ]

    pruned = ctx._prune_tool_outputs(messages)
    assert pruned[1].tool_result.content == "small-output"
    assert pruned[3].tool_result.content == "keep-output"
    assert pruned[5].tool_result.content == "keep-output-2"


def test_repair_incomplete_tool_calls_adds_synthetic_results():
    """修复不完整的工具调用：为缺失 tool_result 的 tool_call 追加合成结果"""
    ctx = DefaultContext(empty_chat_fn, config=ContextConfig(max_tokens=100))
    messages = [
        Message(role=MessageRole.USER, content="u1"),
        _tool_call_msg("assistant call", "t1"),
        Message(role=MessageRole.USER, content="u2"),
    ]

    repaired = ctx._repair_incomplete_tool_calls(messages)

    # 应追加一条合成的 tool_result（在 u2 之后，因为 t1 是最后一个带 tool_calls 的 assistant）
    # 但这里 t1 不是最后一个 assistant with tool_calls（u2 在后面），检查逻辑：
    # 从后往前找最近的 assistant with tool_calls → 是 t1
    # t1 之后的消息中没有 tool_result for t1 → 缺失
    assert len(repaired) == len(messages) + 1
    assert repaired[-1].role == MessageRole.TOOL
    assert repaired[-1].tool_result.tool_call_id == "t1"
    assert repaired[-1].tool_result.is_error is True
    assert "进程中断" in repaired[-1].tool_result.content


@pytest.mark.asyncio
async def test_summarize_messages_inserts_summary_and_keeps_recent_turns():
    fake_chat = FakeChat(["SUM"])
    ctx = DefaultContext(fake_chat, config=ContextConfig(max_tokens=100))

    messages = [
        Message(role=MessageRole.USER, content="u1"),
        Message(role=MessageRole.ASSISTANT, content="a1"),
        Message(role=MessageRole.USER, content="u2"),
        Message(role=MessageRole.ASSISTANT, content="a2"),
    ]

    summarized = await ctx._summarize_messages(
        messages=messages,
        turns=1,
        min_msgs=1,
        strategy="summarize_weak",
    )

    assert summarized[0].role == MessageRole.USER
    assert "<context-summary>" in summarized[0].content
    assert summarized[1].role == MessageRole.ASSISTANT
    assert summarized[-2].content == "u2"
    assert summarized[-1].content == "a2"

    assert len(fake_chat.calls) == 1
    assert any(
        sp.content == ctx._summary_prompt for sp in fake_chat.calls[0]["system_prompts"]
    )


@pytest.mark.asyncio
async def test_summarize_all_returns_only_summary_and_confirmation():
    fake_chat = FakeChat(["ALL"])
    ctx = DefaultContext(fake_chat, config=ContextConfig(max_tokens=100))

    messages = [
        Message(role=MessageRole.USER, content="u1"),
        Message(role=MessageRole.ASSISTANT, content="a1"),
        Message(role=MessageRole.USER, content="u2"),
    ]

    summarized = await ctx._summarize_all(messages)

    assert len(summarized) == 2
    assert summarized[0].role == MessageRole.USER
    assert "<context-summary>" in summarized[0].content
    assert summarized[1].role == MessageRole.ASSISTANT


@pytest.mark.asyncio
async def test_generate_summary_uses_task_prompt_and_history():
    fake_chat = FakeChat(["OK"])
    ctx = DefaultContext(fake_chat, config=ContextConfig(max_tokens=100))
    messages = [
        Message(role=MessageRole.USER, content="Do task A"),
        Message(role=MessageRole.ASSISTANT, content="OK"),
    ]

    summary = await ctx._generate_summary(messages)

    assert summary == "OK"
    assert len(fake_chat.calls) == 1
    call = fake_chat.calls[0]
    assert any(
        sp.content == ctx._summary_prompt for sp in call["system_prompts"]
    )
    assert any(
        "任务摘要" in sp.content for sp in call["system_prompts"]
    )
    assert "对话历史" in call["messages"][0].content


@pytest.mark.asyncio
async def test_compress_if_needed_uses_prune_strategy_when_threshold_met(monkeypatch):
    config = ContextConfig(max_tokens=100, threshold=0.5, prune_threshold=0.5)
    ctx = DefaultContext(empty_chat_fn, config=config)
    ctx._messages = [
        Message(role=MessageRole.USER, content="u1"),
        Message(role=MessageRole.ASSISTANT, content="a1"),
    ]

    async def fake_count(_messages):
        return 10

    monkeypatch.setattr(ctx, "_prune_tool_outputs", lambda msgs: msgs[:1])
    monkeypatch.setattr(ctx, "_count_tokens_for_messages", fake_count)

    did_compress = await ctx.compress_if_needed(force=True)

    assert did_compress is True
    assert len(ctx._messages) == 1
    assert ctx._token_count == 10
