"""Tests for BaseMCPAgent using simulated actions."""

from __future__ import annotations

from typing import Any, ClassVar
from unittest.mock import MagicMock

# Import AsyncMock from unittest.mock if available (Python 3.8+)
try:
    from unittest.mock import AsyncMock
except ImportError:
    # Fallback for older Python versions
    from unittest.mock import MagicMock as AsyncMock

import pytest
from mcp import types

from hud.agents import MCPAgent
from hud.datasets import Task
from hud.tools.executors.base import BaseExecutor
from hud.types import AgentResponse, MCPToolCall, MCPToolResult, Trace


class MockMCPAgent(MCPAgent):
    """Concrete implementation of BaseMCPAgent for testing."""

    metadata: ClassVar[dict[str, Any]] = {}  # Optional metadata for MCP config

    def __init__(self, mcp_client: Any = None, **kwargs: Any) -> None:
        if mcp_client is None:
            # Create a mock client if none provided
            mcp_client = MagicMock()
            mcp_client.get_available_tools = MagicMock(return_value=[])
            mcp_client.initialize = AsyncMock()
            mcp_client.list_tools = AsyncMock(return_value=[])
            mcp_client.mcp_config = {"test_server": {"url": "http://localhost"}}
        super().__init__(mcp_client=mcp_client, **kwargs)
        self.executor = BaseExecutor()  # Use simulated executor
        self._messages = []

    async def run(self, task: Task) -> list[dict[str, Any]]:
        """Mock run method."""
        return self._messages

    async def create_initial_messages(
        self, prompt: str, initial_screenshot: bool = False
    ) -> list[dict[str, Any]]:
        """Mock create initial messages."""
        messages = [{"role": "user", "content": prompt}]
        if initial_screenshot:
            messages.append({"role": "assistant", "content": "Screenshot: mock_screenshot"})
        return messages

    async def get_response(self, messages: list[dict[str, Any]]) -> AgentResponse:
        """Mock get response."""
        return AgentResponse(content="Mock response", tool_calls=[], done=True)

    async def format_tool_results(
        self, tool_calls: list[MCPToolCall], tool_results: list[MCPToolResult]
    ) -> list[dict[str, Any]]:
        """Mock format tool results."""
        formatted = []
        for tool_call, result in zip(tool_calls, tool_results):
            formatted.append({"role": "tool", "name": tool_call.name, "content": str(result)})
        return formatted

    async def create_user_message(self, text: str) -> Any:
        """Mock create user message."""
        return {"role": "user", "content": text}

    async def get_system_messages(self) -> list[Any]:
        """Mock get system messages."""
        return []

    async def format_blocks(self, blocks: list[types.ContentBlock]) -> list[Any]:
        """Mock format blocks."""
        formatted = []
        for block in blocks:
            if isinstance(block, types.TextContent):
                formatted.append({"type": "text", "text": block.text})
            elif isinstance(block, types.ImageContent):
                formatted.append({"type": "image", "data": block.data})
            elif hasattr(block, "type"):
                formatted.append({"type": getattr(block, "type", "unknown")})
        return formatted


class TestBaseMCPAgent:
    """Tests for BaseMCPAgent with simulated actions."""

    def test_init_defaults(self):
        """Test initialization with default values."""
        agent = MockMCPAgent()

        assert agent.mcp_client is not None
        assert agent.allowed_tools is None
        assert agent.disallowed_tools == []
        assert agent.initial_screenshot is True
        assert agent.system_prompt is not None  # Default system prompt is set

    def test_init_with_params(self):
        """Test initialization with custom parameters."""
        client = MagicMock()
        agent = MockMCPAgent(
            mcp_client=client,
            allowed_tools=["tool1", "tool2"],
            disallowed_tools=["bad_tool"],
            initial_screenshot=True,
            system_prompt="Custom prompt",
        )

        assert agent.mcp_client == client
        assert agent.allowed_tools == ["tool1", "tool2"]
        assert agent.disallowed_tools == ["bad_tool"]
        assert agent.initial_screenshot is True
        assert agent.system_prompt == "Custom prompt"

    @pytest.mark.asyncio
    async def test_init_no_client_no_task(self):
        """Test initialize fails without client and without task."""

        # Create a minimal concrete implementation to test the ValueError
        class TestAgent(MCPAgent):
            async def create_initial_messages(
                self, prompt: str, initial_screenshot: bool = False
            ) -> list[dict[str, Any]]:
                return []

            async def format_tool_results(
                self, tool_calls: list[MCPToolCall], tool_results: list[MCPToolResult]
            ) -> list[dict[str, Any]]:
                return []

            async def get_response(self, messages: list[dict[str, Any]]) -> dict[str, Any]:
                return {"content": "test"}

            async def get_system_messages(self) -> list[Any]:
                return []

            async def format_blocks(self, blocks: list[types.ContentBlock]) -> list[Any]:
                return []

        # Agent can be created with None client
        agent = TestAgent(mcp_client=None)

        # But initialize should fail without client or task
        with pytest.raises(ValueError, match="No MCPClient"):
            await agent.initialize()

    @pytest.mark.asyncio
    async def test_initialize_with_sessions(self):
        """Test initialize with existing sessions."""
        agent = MockMCPAgent()

        # Create proper async mock for session
        mock_session = MagicMock()

        # Set up the connector and client_session structure
        mock_session.connector = MagicMock()
        mock_session.connector.client_session = MagicMock()

        # Mock list_tools on the client_session
        async def mock_list_tools():
            return types.ListToolsResult(
                tools=[
                    types.Tool(name="tool1", description="Tool 1", inputSchema={"type": "object"}),
                    types.Tool(name="tool2", description="Tool 2", inputSchema={"type": "object"}),
                    types.Tool(
                        name="setup", description="Setup tool", inputSchema={"type": "object"}
                    ),
                ]
            )

        mock_session.connector.client_session.list_tools = mock_list_tools

        assert agent.mcp_client is not None

        # Mock the list_tools method on mcp_client to return the tools
        agent.mcp_client.list_tools = AsyncMock(
            return_value=[
                types.Tool(name="tool1", description="Tool 1", inputSchema={"type": "object"}),
                types.Tool(name="tool2", description="Tool 2", inputSchema={"type": "object"}),
                types.Tool(name="setup", description="Setup tool", inputSchema={"type": "object"}),
            ]
        )

        await agent.initialize()

        # Check available tools were populated (excludes lifecycle tools)
        tools = agent.get_available_tools()
        assert len(tools) == 3  # All tools (setup is not in default lifecycle tools)

        # Ensure names exist in available tools
        names = {t.name for t in tools}
        assert {"tool1", "tool2", "setup"} <= names

    @pytest.mark.asyncio
    async def test_initialize_with_filtering(self):
        """Test initialize with tool filtering."""
        agent = MockMCPAgent(allowed_tools=["tool1"], disallowed_tools=["tool3"])

        # Create proper async mock for session
        mock_session = MagicMock()

        # Set up the connector and client_session structure
        mock_session.connector = MagicMock()
        mock_session.connector.client_session = MagicMock()

        async def mock_list_tools():
            return types.ListToolsResult(
                tools=[
                    types.Tool(name="tool1", description="Tool 1", inputSchema={"type": "object"}),
                    types.Tool(name="tool2", description="Tool 2", inputSchema={"type": "object"}),
                    types.Tool(name="tool3", description="Tool 3", inputSchema={"type": "object"}),
                    types.Tool(name="setup", description="Setup", inputSchema={"type": "object"}),
                ]
            )

        mock_session.connector.client_session.list_tools = mock_list_tools

        assert agent.mcp_client is not None

        # Mock the list_tools method on mcp_client to return the tools
        agent.mcp_client.list_tools = AsyncMock(
            return_value=[
                types.Tool(name="tool1", description="Tool 1", inputSchema={"type": "object"}),
                types.Tool(name="tool2", description="Tool 2", inputSchema={"type": "object"}),
                types.Tool(name="tool3", description="Tool 3", inputSchema={"type": "object"}),
                types.Tool(name="setup", description="Setup", inputSchema={"type": "object"}),
            ]
        )

        await agent.initialize()

        # Check filtering worked - get_available_tools excludes lifecycle tools
        tools = agent.get_available_tools()
        tool_names = [t.name for t in tools]
        assert len(tools) == 1  # Only tool1 (tool2 and tool3 are filtered out)
        assert "tool1" in tool_names
        assert "setup" not in tool_names  # Lifecycle tool excluded from available tools
        assert "tool2" not in tool_names  # Not in allowed list
        assert "tool3" not in tool_names  # In disallowed list

    @pytest.mark.asyncio
    async def test_call_tool_success(self):
        """Test successful tool call."""
        agent = MockMCPAgent()

        # Initialize with a tool
        mock_session = MagicMock()
        mock_session.connector = MagicMock()
        mock_session.connector.client_session = MagicMock()

        async def mock_list_tools():
            return types.ListToolsResult(
                tools=[
                    types.Tool(name="test_tool", description="Test", inputSchema={"type": "object"})
                ]
            )

        mock_session.connector.client_session.list_tools = mock_list_tools

        # Mock the call_tool method on the client session
        mock_result = types.CallToolResult(
            content=[types.TextContent(type="text", text="Tool result")], isError=False
        )

        async def mock_call_tool(name, args):
            return mock_result

        mock_session.connector.client_session.call_tool = mock_call_tool

        assert agent.mcp_client is not None

        # Mock the client's call_tool method directly
        agent.mcp_client.call_tool = AsyncMock(return_value=mock_result)

        # Mock the list_tools method to return the test tool
        agent.mcp_client.list_tools = AsyncMock(
            return_value=[
                types.Tool(name="test_tool", description="Test", inputSchema={"type": "object"})
            ]
        )

        await agent.initialize()

        # Call the tool
        tool_call = MCPToolCall(name="test_tool", arguments={"param": "value"})
        results = await agent.call_tools(tool_call)

        assert len(results) == 1
        assert results[0] == mock_result
        assert not results[0].isError

    @pytest.mark.asyncio
    async def test_call_tool_not_found(self):
        """Test calling non-existent tool."""
        agent = MockMCPAgent()

        # Initialize without tools
        mock_session = MagicMock()

        async def mock_list_tools():
            return types.ListToolsResult(tools=[])

        mock_session.list_tools = mock_list_tools
        assert agent.mcp_client is not None

        await agent.initialize()

        # Try to call unknown tool - call_tools doesn't raise for unknown tools
        tool_call = MCPToolCall(name="unknown_tool", arguments={})
        await agent.call_tools(tool_call)

    @pytest.mark.asyncio
    async def test_call_tool_no_name(self):
        """Test calling tool without name."""
        # MCPToolCall accepts empty names
        agent = MockMCPAgent()
        tool_call = MCPToolCall(name="", arguments={})

        # call_tools doesn't validate empty names, it will return error
        await agent.call_tools(tool_call)

    def test_get_tool_schemas(self):
        """Test getting tool schemas."""
        agent = MockMCPAgent()

        agent._available_tools = [
            types.Tool(name="tool1", description="Tool 1", inputSchema={"type": "object"}),
            types.Tool(name="setup", description="Setup", inputSchema={"type": "object"}),
        ]

        schemas = agent.get_tool_schemas()

        # Should include non-lifecycle tools
        assert len(schemas) == 1
        assert schemas[0]["name"] == "tool1"

    def test_get_tools_by_server(self):
        """Test getting tools grouped by server."""
        agent = MockMCPAgent()

        # Set up tools from different servers
        tool1 = types.Tool(name="tool1", description="Tool 1", inputSchema={"type": "object"})
        tool2 = types.Tool(name="tool2", description="Tool 2", inputSchema={"type": "object"})

        agent._available_tools = [tool1, tool2]
        tools = agent.get_available_tools()
        assert {t.name for t in tools} == {"tool1", "tool2"}

    @pytest.mark.asyncio
    async def test_executor_integration(self):
        """Test integration with BaseExecutor for simulated actions."""
        agent = MockMCPAgent()

        # Test various executor actions
        click_result = await agent.executor.click(100, 200, take_screenshot=False)
        assert click_result.output is not None
        assert "[SIMULATED] Click at (100, 200)" in click_result.output

        type_result = await agent.executor.write("Test input", take_screenshot=False)
        assert type_result.output is not None
        assert "[SIMULATED] Type 'Test input'" in type_result.output

        scroll_result = await agent.executor.scroll(x=50, y=50, scroll_y=5, take_screenshot=False)
        assert scroll_result.output is not None
        assert "[SIMULATED] Scroll" in scroll_result.output

        # Test screenshot
        screenshot = await agent.executor.screenshot()
        assert isinstance(screenshot, str)
        assert screenshot.startswith("iVBORw0KGgo")  # PNG header


class MockAgentExtended(MCPAgent):
    """Mock agent for testing with predefined responses."""

    metadata: ClassVar[dict[str, Any]] = {}  # Optional metadata for MCP config

    def __init__(self, responses=None, **kwargs):
        super().__init__(**kwargs)
        self.responses = responses or []
        self.call_count = 0

    async def create_initial_messages(
        self, prompt: str, initial_screenshot: bool = False
    ) -> list[dict[str, Any]]:
        """Create initial messages."""
        messages = [{"role": "user", "content": prompt}]
        if initial_screenshot:
            # capture_screenshot doesn't exist, just mock it
            screenshot = "mock_screenshot_data"
            messages.append({"role": "assistant", "content": f"Screenshot: {screenshot}"})
        return messages

    async def get_response(self, messages: list[dict[str, Any]]) -> AgentResponse:
        """Return predefined responses - must be async."""
        if self.call_count < len(self.responses):
            response_dict = self.responses[self.call_count]
            self.call_count += 1
            # Convert dict to AgentResponse
            return AgentResponse(
                content=response_dict.get("content", ""),
                tool_calls=response_dict.get("tool_calls", []),
                done=response_dict.get("done", not bool(response_dict.get("tool_calls"))),
            )
        return AgentResponse(content="Done", tool_calls=[], done=True)

    async def format_tool_results(
        self, tool_calls: list[MCPToolCall], tool_results: list[MCPToolResult]
    ) -> list[dict[str, Any]]:
        """Format tool results."""
        formatted = []
        for tool_call, result in zip(tool_calls, tool_results):
            formatted.append({"role": "tool", "name": tool_call.name, "content": str(result)})
        return formatted

    async def create_user_message(self, text: str) -> Any:
        """Create user message."""
        return {"role": "user", "content": text}

    async def get_system_messages(self) -> list[Any]:
        """Mock get system messages."""
        return []

    async def format_blocks(self, blocks: list[types.ContentBlock]) -> list[Any]:
        """Mock format blocks."""
        formatted = []
        for block in blocks:
            if isinstance(block, types.TextContent):
                formatted.append({"type": "text", "text": block.text})
            elif isinstance(block, types.ImageContent):
                formatted.append({"type": "image", "data": block.data})
            elif hasattr(block, "type"):
                formatted.append({"type": getattr(block, "type", "unknown")})
        return formatted


class TestMCPAgentExtended:
    """Extended tests for MCPAgent."""

    @pytest.fixture
    def mock_client(self):
        """Create a mock MCP client."""
        client = MagicMock()
        client.get_all_active_sessions = MagicMock(return_value={})
        client.initialize = AsyncMock()
        client.list_tools = AsyncMock(return_value=[])
        client.call_tool = AsyncMock(
            return_value=types.CallToolResult(
                content=[types.TextContent(type="text", text="Success")],
                isError=False,
            )
        )
        return client

    @pytest.fixture
    def agent_with_tools(self, mock_client):
        """Create agent with mock tools."""
        mock_client.list_tools = AsyncMock(
            return_value=[
                types.Tool(name="screenshot", description="Take screenshot", inputSchema={}),
                types.Tool(name="click", description="Click at coordinates", inputSchema={}),
                types.Tool(name="type", description="Type text", inputSchema={}),
                types.Tool(name="bad_tool", description="A tool that fails", inputSchema={}),
            ]
        )
        return MockAgentExtended(mcp_client=mock_client)

    @pytest.mark.asyncio
    async def test_run_with_task_object(self, agent_with_tools):
        """Test running agent with Task object."""
        from hud.types import MCPToolResult

        task = Task(
            id="test_task",
            prompt="Click the button",
            mcp_config={"test_server": {"url": "http://localhost:8080"}},
            setup_tool={"name": "navigate", "arguments": {"url": "https://example.com"}},  # type: ignore[arg-type]
            evaluate_tool={"name": "check_result", "arguments": {}},  # type: ignore[arg-type]
        )

        # Set up responses
        agent_with_tools.responses = [
            {
                "role": "assistant",
                "content": "I'll click the button",
                "tool_calls": [MCPToolCall(name="click", arguments={"x": 100, "y": 200})],
            }
        ]

        # Mock the evaluation to return a reward
        agent_with_tools.mcp_client.call_tool = AsyncMock(
            side_effect=[
                # Setup tool
                MCPToolResult(
                    content=[types.TextContent(type="text", text="Navigated")],
                    isError=False,
                ),
                # Click tool
                MCPToolResult(
                    content=[types.TextContent(type="text", text="Clicked")],
                    isError=False,
                ),
                # Evaluate tool with reward
                MCPToolResult(
                    content=[types.TextContent(type="text", text="Success")],
                    isError=False,
                    structuredContent={"reward": 1.0},
                ),
            ]
        )

        result = await agent_with_tools.run(task)

        assert isinstance(result, Trace)
        assert result.reward == 1.0
        assert not result.isError
        assert result.done

    @pytest.mark.asyncio
    async def test_run_with_setup_error(self, agent_with_tools):
        """Test task execution with setup phase error."""
        from hud.types import MCPToolResult

        task = Task(
            id="test_task",
            prompt="Do something",
            mcp_config={"test_server": {"url": "http://localhost:8080"}},
            setup_tool={"name": "bad_setup", "arguments": {}},  # type: ignore[arg-type]
        )

        # Mock setup tool to fail
        agent_with_tools.mcp_client.call_tool = AsyncMock(
            return_value=MCPToolResult(
                content=[types.TextContent(type="text", text="Setup failed")],
                isError=True,
            )
        )

        result = await agent_with_tools.run(task)

        assert isinstance(result, Trace)
        assert result.isError
        # Error content is the string representation of the MCPToolResult list
        assert result.content is not None
        assert "Setup failed" in result.content
        assert "MCPToolResult" in result.content

    @pytest.mark.asyncio
    async def test_run_with_multiple_setup_tools(self, agent_with_tools):
        """Test task with multiple setup tools."""

        task = Task(
            id="test_task",
            prompt="Test multiple setup",
            mcp_config={"test_server": {"url": "http://localhost:8080"}},
            setup_tool=[
                MCPToolCall(name="setup1", arguments={}),
                MCPToolCall(name="setup2", arguments={}),
            ],
        )

        agent_with_tools.responses = [{"role": "assistant", "content": "Done", "tool_calls": []}]

        setup_calls = []
        agent_with_tools.mcp_client.call_tool = AsyncMock(
            side_effect=lambda tool_call: setup_calls.append(tool_call)
            or MCPToolResult(
                content=[types.TextContent(type="text", text=f"{tool_call.name} done")],
                isError=False,
            )
        )

        result = await agent_with_tools.run(task)

        # Check that the tool names match
        setup_names = [call.name for call in setup_calls]
        assert "setup1" in setup_names
        assert "setup2" in setup_names
        assert not result.isError

    @pytest.mark.asyncio
    async def test_allowed_tools_filtering(self, mock_client):
        """Test that allowed_tools filters available tools."""
        mock_client.list_tools = AsyncMock(
            return_value=[
                types.Tool(name="tool1", description="Tool 1", inputSchema={}),
                types.Tool(name="tool2", description="Tool 2", inputSchema={}),
                types.Tool(name="tool3", description="Tool 3", inputSchema={}),
            ]
        )

        agent = MockAgentExtended(mcp_client=mock_client, allowed_tools=["tool1", "tool3"])
        await agent.initialize("test")

        available_names = [tool.name for tool in agent.get_available_tools()]
        assert "tool1" in available_names
        assert "tool3" in available_names
        assert "tool2" not in available_names

    @pytest.mark.asyncio
    async def test_disallowed_tools_filtering(self, mock_client):
        """Test that disallowed_tools filters available tools."""
        mock_client.list_tools = AsyncMock(
            return_value=[
                types.Tool(name="tool1", description="Tool 1", inputSchema={}),
                types.Tool(name="tool2", description="Tool 2", inputSchema={}),
                types.Tool(name="tool3", description="Tool 3", inputSchema={}),
            ]
        )

        agent = MockAgentExtended(mcp_client=mock_client, disallowed_tools=["tool2"])
        await agent.initialize("test")

        available_names = [tool.name for tool in agent.get_available_tools()]
        assert "tool1" in available_names
        assert "tool3" in available_names
        assert "tool2" not in available_names

    @pytest.mark.asyncio
    async def test_lifecycle_tools(self, mock_client):
        """Test lifecycle tools are called in run_prompt."""
        # Lifecycle tools are specified by name, not as objects
        agent = MockAgentExtended(
            mcp_client=mock_client,
            responses=[{"role": "assistant", "content": "Done", "tool_calls": []}],
        )

        # Add screenshot tool to available tools
        mock_client.list_tools = AsyncMock(
            return_value=[
                types.Tool(name="screenshot", description="Take screenshot", inputSchema={})
            ]
        )

        # Initialize to make tools available
        await agent.initialize()

        result = await agent.run("Test lifecycle", max_steps=1)
        assert not result.isError

    # This test is commented out as screenshot history management may have changed
    # @pytest.mark.asyncio
    # async def test_screenshot_history_management(self, agent_with_tools):
    #     """Test screenshot history is maintained."""
    #     agent_with_tools.initial_screenshot = True

    #     # Set up responses with tool calls
    #     agent_with_tools.responses = [
    #         {
    #             "role": "assistant",
    #             "content": "Action 1",
    #             "tool_calls": [MCPToolCall(name="click", arguments={"x": 1, "y": 1})],
    #         },
    #         {
    #             "role": "assistant",
    #             "content": "Action 2",
    #             "tool_calls": [MCPToolCall(name="click", arguments={"x": 2, "y": 2})],
    #         },
    #         {
    #             "role": "assistant",
    #             "content": "Action 3",
    #             "tool_calls": [MCPToolCall(name="click", arguments={"x": 3, "y": 3})],
    #         },
    #     ]

    #     await agent_with_tools.run("Test screenshots", max_steps=3)

    #     # Should have screenshots in history
    #     assert len(agent_with_tools.screenshot_history) > 0

    @pytest.mark.asyncio
    async def test_run_with_invalid_prompt_type(self, agent_with_tools):
        """Test run with invalid prompt type raises TypeError."""
        with pytest.raises(TypeError, match="prompt_or_task must be str or Task"):
            await agent_with_tools.run(123)  # Invalid type

    @pytest.mark.asyncio
    async def test_evaluate_phase_with_multiple_tools(self, agent_with_tools):
        """Test evaluation phase with multiple evaluation tools."""
        from hud.types import MCPToolResult

        task = Task(
            id="test_task",
            prompt="Test evaluation",
            mcp_config={"test_server": {"url": "http://localhost:8080"}},
            evaluate_tool=[
                MCPToolCall(name="eval1", arguments={}),
                MCPToolCall(name="eval2", arguments={"reward": True}),
            ],
        )

        agent_with_tools.responses = [{"role": "assistant", "content": "Done", "tool_calls": []}]

        eval_calls = []
        agent_with_tools.mcp_client.call_tool = AsyncMock(
            side_effect=lambda tool_call: eval_calls.append(tool_call)
            or MCPToolResult(
                content=[types.TextContent(type="text", text=f"{tool_call.name} result")],
                isError=False,
                structuredContent={"reward": 0.5} if tool_call.name == "eval1" else {"reward": 1.0},
            )
        )

        result = await agent_with_tools.run(task)

        # Check that the tool names match
        eval_names = [call.name for call in eval_calls]
        assert "eval1" in eval_names
        assert "eval2" in eval_names
        assert result.reward == 0.5  # From eval1 (first evaluation tool)

    @pytest.mark.asyncio
    async def test_trace_population_on_error(self, agent_with_tools):
        """Test that trace is populated on task execution error."""

        task = Task(
            id="test_task",
            prompt="Test error",
            mcp_config={"test_server": {"url": "http://localhost:8080"}},
            setup_tool={"name": "failing_setup", "arguments": {}},  # type: ignore[arg-type]
        )

        # Make setup fail with exception
        agent_with_tools.mcp_client.call_tool = AsyncMock(side_effect=Exception("Setup explosion"))

        result = await agent_with_tools.run(task)

        assert result.isError
        # Error content is the string representation of the MCPToolResult list
        assert "Setup explosion" in result.content
        assert "MCPToolResult" in result.content
        assert result.done
