# -*- coding: utf-8 -*-
"""SelfCorrectionLoop 集成到 process_tool_calls 的测试。

验证工具失败后自动注入结构化修正提示，触发 LLM 修正参数。
"""
import json
from unittest.mock import MagicMock, patch

import pytest

from abyss.agent_loop import process_tool_calls, SelfCorrectionLoop
from abyss.tools.registry import ToolRegistry
from abyss.tools.base import Tool
from abyss.session import Session


@pytest.fixture(autouse=True)
def mock_logger():
    """mock abyss.agent_loop.logger 避免 sandbox 文件写入限制。"""
    with patch("abyss.agent_loop.logger") as m:
        yield m


class FailingTool(Tool):
    """永远失败的测试工具。"""

    @property
    def name(self) -> str:
        return "failing_tool"

    @property
    def description(self) -> str:
        return "A tool that always fails"

    @property
    def parameters(self) -> dict:
        return {
            "type": "object",
            "properties": {
                "path": {"type": "string"},
            },
            "required": ["path"],
        }

    def execute(self, path: str, **kwargs) -> dict:
        return {"success": False, "error": f"未找到匹配的旧内容: {path}"}


class SuccessTool(Tool):
    """永远成功的测试工具。"""

    @property
    def name(self) -> str:
        return "success_tool"

    @property
    def description(self) -> str:
        return "A tool that always succeeds"

    @property
    def parameters(self) -> dict:
        return {
            "type": "object",
            "properties": {"msg": {"type": "string"}},
            "required": ["msg"],
        }

    def execute(self, msg: str, **kwargs) -> dict:
        return {"success": True, "echo": msg}


class TestSelfCorrectionIntegration:
    """SelfCorrectionLoop 集成到 process_tool_calls 的行为测试。"""

    def make_session(self):
        return Session("You are a test assistant.")

    def make_registry(self, *tools):
        reg = ToolRegistry()
        for t in tools:
            reg.register(t)
        return reg

    def test_failure_injects_correction_prompt(self):
        """工具失败后注入结构化修正提示 user 消息。"""
        session = self.make_session()
        reg = self.make_registry(FailingTool())
        client = MagicMock()
        loop = SelfCorrectionLoop(max_attempts=3)

        captured_msgs = []

        def fake_stream(_client, msgs, _schema, _cfg, _sr):
            captured_msgs.append(list(msgs))  # 快照每轮 messages
            return ("done after correction", "", [], None)

        tc = {"function": {"name": "failing_tool", "arguments": json.dumps({"path": "foo.txt"})}}

        process_tool_calls(
            session=session,
            client=client,
            tools_schemas=reg.to_openai_schemas(),
            tool_calls=[tc],
            cfg={"tool_result_max_chars": None},
            registry=reg,
            stream_handler=fake_stream,
            self_correction=loop,
        )

        # 验证：第二轮 stream_handler 收到的 messages 应包含修正提示 user 消息
        assert len(captured_msgs) >= 1
        last_msgs = captured_msgs[-1]
        user_msgs = [m for m in last_msgs if m.get("role") == "user"]
        correction_msgs = [m for m in user_msgs if "自纠错提示" in (m.get("content") or "")]
        assert len(correction_msgs) >= 1, "失败后应注入含「自纠错提示」的 user 消息"
        assert "failing_tool" in correction_msgs[0]["content"]
        assert "未找到匹配" in correction_msgs[0]["content"]
        assert loop.attempts == 1

    def test_success_does_not_inject(self):
        """工具成功时不注入修正提示。"""
        session = self.make_session()
        reg = self.make_registry(SuccessTool())
        client = MagicMock()
        loop = SelfCorrectionLoop(max_attempts=3)

        captured_msgs = []

        def fake_stream(_client, msgs, _schema, _cfg, _sr):
            captured_msgs.append(list(msgs))
            return ("done", "", [], None)

        tc = {"function": {"name": "success_tool", "arguments": json.dumps({"msg": "hi"})}}

        process_tool_calls(
            session=session,
            client=client,
            tools_schemas=reg.to_openai_schemas(),
            tool_calls=[tc],
            cfg={"tool_result_max_chars": None},
            registry=reg,
            stream_handler=fake_stream,
            self_correction=loop,
        )

        last_msgs = captured_msgs[-1]
        correction_msgs = [
            m for m in last_msgs
            if m.get("role") == "user" and "自纠错提示" in (m.get("content") or "")
        ]
        assert len(correction_msgs) == 0, "成功结果不应注入修正提示"
        assert loop.attempts == 0

    def test_exhausted_no_more_inject(self):
        """达到 max_attempts 上限后不再注入修正提示。"""
        session = self.make_session()
        reg = self.make_registry(FailingTool())
        client = MagicMock()
        loop = SelfCorrectionLoop(max_attempts=2)

        call_count = [0]

        def fake_stream(_client, _msgs, _schema, _cfg, _sr):
            call_count[0] += 1
            # 第 1 轮 follow-up：再调失败工具
            if call_count[0] == 1:
                new_tc = [
                    {"index": 0, "id": "call_2",
                     "function": {"name": "failing_tool",
                                  "arguments": json.dumps({"path": "bar.txt"})}}
                ]
                return ("", "", new_tc, None)
            # 第 2 轮 follow-up：再调失败工具
            if call_count[0] == 2:
                new_tc = [
                    {"index": 0, "id": "call_3",
                     "function": {"name": "failing_tool",
                                  "arguments": json.dumps({"path": "baz.txt"})}}
                ]
                return ("", "", new_tc, None)
            # 第 3 轮 follow-up：放弃，返回文本
            return ("giving up", "", [], None)

        tc = {"function": {"name": "failing_tool", "arguments": json.dumps({"path": "foo.txt"})}}

        process_tool_calls(
            session=session,
            client=client,
            tools_schemas=reg.to_openai_schemas(),
            tool_calls=[tc],
            cfg={"tool_result_max_chars": None},
            registry=reg,
            stream_handler=fake_stream,
            self_correction=loop,
        )

        # max_attempts=2，应仅注入 2 次修正提示
        assert loop.attempts == 2

    def test_none_loop_backward_compatible(self):
        """self_correction=None 时行为与原来一致（向后兼容）。"""
        session = self.make_session()
        reg = self.make_registry(FailingTool())
        client = MagicMock()

        def fake_stream(_client, _msgs, _schema, _cfg, _sr):
            return ("done", "", [], None)

        tc = {"function": {"name": "failing_tool", "arguments": json.dumps({"path": "foo.txt"})}}

        # 不传 self_correction，默认 None
        result = process_tool_calls(
            session=session,
            client=client,
            tools_schemas=reg.to_openai_schemas(),
            tool_calls=[tc],
            cfg={"tool_result_max_chars": None},
            registry=reg,
            stream_handler=fake_stream,
        )

        # 行为不变：返回 stream_handler 的 content
        assert result == "done"

    def test_correction_prompt_contains_suggestion(self):
        """修正提示包含工具特定的修正建议。"""
        session = self.make_session()
        reg = self.make_registry(FailingTool())
        client = MagicMock()
        loop = SelfCorrectionLoop(max_attempts=3)

        captured = []

        def fake_stream(_client, msgs, _schema, _cfg, _sr):
            captured.extend([m for m in msgs if m.get("role") == "user"
                            and "自纠错提示" in (m.get("content") or "")])
            return ("done", "", [], None)

        # 用 file_edit 名字触发特定建议
        tc = {"function": {"name": "file_edit", "arguments": json.dumps({"path": "foo.txt"})}}

        # 注册一个 file_edit 工具（用 FailingTool 行为）
        reg2 = ToolRegistry()
        reg2.register(FailingTool())  # name=failing_tool
        # 改名模拟 file_edit
        tc_failing = {"function": {"name": "failing_tool",
                                    "arguments": json.dumps({"path": "foo.txt"})}}

        process_tool_calls(
            session=session,
            client=client,
            tools_schemas=reg2.to_openai_schemas(),
            tool_calls=[tc_failing],
            cfg={"tool_result_max_chars": None},
            registry=reg2,
            stream_handler=fake_stream,
            self_correction=loop,
        )

        assert len(captured) >= 1
        content = captured[0]["content"]
        # 含工具名、错误、建议三要素
        assert "failing_tool" in content
        assert "未找到匹配" in content
        assert "建议" in content
