"""
测试 main_agent.py — _save_iteration_messages 对 fetch_url 结果清洗后持久化
"""
import sys
import json
import tempfile
import os
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch

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


class MockSessionManager:
    def __init__(self):
        self.current_messages = []
        self.current_session_id = "test_session"
        self.sessions_index = {}
        self.sessions_dir = Path(tempfile.mkdtemp())

    async def _save_session(self):
        pass

    def _save_sessions_index(self):
        pass

    async def _save_current_session_id(self):
        pass

    def _get_total_tokens(self):
        return 0


def test_fetch_url_content_cleaned():
    """验证 fetch_url 工具消息的 HTML 内容在持久化时被清洗"""
    # Arrange
    html_content = {
        "success": True,
        "url": "https://pytorch.org/docs/nn.html",
        "status_code": 200,
        "content_type": "text/html",
        "content_length": 5000,
        "content": (
            "<!DOCTYPE html>\n<html lang='en'>\n<head>\n"
            "<link href='css_1.css' rel='stylesheet'/>"
            "<link href='css_2.css' rel='stylesheet'/>"
            "<script>console.log('hi')</script>\n"
            "<style>.a{color:red}</style>\n"
            "</head>\n<body>\n"
            "<h1>torch.nn — PyTorch 2.12 documentation</h1>\n"
            "<p>Containers, Convolution Layers, Pooling layers.</p>\n"
            "</body></html>"
        ),
    }

    new_msgs = [
        {
            "role": "assistant",
            "content": "",
            "tool_calls": [
                {
                    "id": "call_001",
                    "type": "function",
                    "function": {"name": "fetch_url", "arguments": "{}"},
                }
            ],
        },
        {
            "role": "tool",
            "tool_call_id": "call_001",
            "content": str(html_content),
        },
    ]

    from aacode.core.main_agent import MainAgent
    from aacode.tools.web_tools import WebTools

    mock_session_mgr = MockSessionManager()

    agent = MagicMock(spec=MainAgent)
    agent.session_manager = mock_session_mgr
    agent.web_tools = WebTools.__new__(WebTools)  # 提供 _clean_html 静态方法
    agent._save_iteration_messages = MainAgent._save_iteration_messages.__get__(agent, MainAgent)

    import asyncio
    asyncio.run(agent._save_iteration_messages(new_msgs))

    assert len(mock_session_mgr.current_messages) == 2
    tool_msg = mock_session_mgr.current_messages[1]
    saved_content = tool_msg.content

    # 不应包含 CSS/JS 垃圾
    assert "css_1.css" not in saved_content
    assert "console.log" not in saved_content
    assert ".a{color:red}" not in saved_content
    assert "<head>" not in saved_content
    assert "<script>" not in saved_content
    assert "<style>" not in saved_content

    # 应包含页面正文
    assert "torch.nn" in saved_content
    assert "PyTorch 2.12" in saved_content
    assert "Containers" in saved_content

    print("  test_fetch_url_content_cleaned: PASS")


def test_non_fetch_url_tool_preserved():
    """非 fetch_url 的工具消息不做清洗"""
    new_msgs = [
        {
            "role": "assistant",
            "content": "",
            "tool_calls": [
                {
                    "id": "call_002",
                    "type": "function",
                    "function": {"name": "run_shell", "arguments": "{}"},
                }
            ],
        },
        {
            "role": "tool",
            "tool_call_id": "call_002",
            "content": "<html>Should be preserved as-is</html>",
        },
    ]

    from aacode.core.main_agent import MainAgent

    mock_session_mgr = MockSessionManager()
    agent = MagicMock(spec=MainAgent)
    agent.session_manager = mock_session_mgr
    agent._save_iteration_messages = MainAgent._save_iteration_messages.__get__(agent, MainAgent)

    import asyncio
    asyncio.run(agent._save_iteration_messages(new_msgs))

    tool_msg = mock_session_mgr.current_messages[1]
    saved_content = tool_msg.content
    assert "<html>" in saved_content
    assert "preserved" in saved_content

    print("  test_non_fetch_url_tool_preserved: PASS")


if __name__ == "__main__":
    test_fetch_url_content_cleaned()
    test_non_fetch_url_tool_preserved()
