import logging
from unittest.mock import Mock

import pytest
from pydantic import BaseModel

from tests.test_obj_factory import get_event_obj
from unique_toolkit.agentic.tools.a2a import SubAgentResponseWatcher
from unique_toolkit.agentic.tools.a2a.manager import A2AManager
from unique_toolkit.agentic.tools.config import (
    ToolBuildConfig,
    ToolIcon,
    ToolSelectionPolicy,
)
from unique_toolkit.agentic.tools.factory import ToolFactory
from unique_toolkit.agentic.tools.mcp.manager import MCPManager
from unique_toolkit.agentic.tools.mcp.models import MCPToolConfig
from unique_toolkit.agentic.tools.mcp.tool_wrapper import MCPToolWrapper
from unique_toolkit.agentic.tools.schemas import BaseToolConfig
from unique_toolkit.agentic.tools.tool import Tool
from unique_toolkit.agentic.tools.tool_manager import ToolManager, ToolManagerConfig
from unique_toolkit.agentic.tools.tool_progress_reporter import ToolProgressReporter
from unique_toolkit.app.schemas import McpServer, McpTool
from unique_toolkit.chat.service import ChatService


class MockParameters(BaseModel):
    pass


class MockInternalSearchTool(Tool[BaseToolConfig]):
    """Mock internal search tool for testing"""

    name = "internal_search"

    def __init__(self, config, event, tool_progress_reporter=None):
        super().__init__(config, event, tool_progress_reporter)

    def tool_description(self):
        from unique_toolkit.language_model.schemas import LanguageModelToolDescription

        return LanguageModelToolDescription(
            name="internal_search",
            description="Internal search tool for testing",
            parameters=MockParameters,
        )

    def tool_description_for_system_prompt(self) -> str:
        return "Internal search tool for searching content"

    def tool_format_information_for_system_prompt(self) -> str:
        return "Use this tool to search for content"

    def evaluation_check_list(self):
        return []

    def get_evaluation_checks_based_on_tool_response(self, tool_response):
        return []

    def get_tool_prompts(self):
        from unique_toolkit.agentic.tools.schemas import ToolPrompts

        return ToolPrompts()

    async def run(self, tool_call):
        from unique_toolkit.agentic.tools.schemas import ToolCallResponse

        return ToolCallResponse(id=tool_call.id, name=tool_call.name, content_chunks=[])


class MockInternalSearchConfig(BaseToolConfig):
    """Mock configuration for internal search tool"""

    pass


class TestMCPManager:
    @pytest.fixture(autouse=True)
    def setup(self):
        """Setup test environment"""
        self.logger = logging.getLogger(__name__)

        # Register mock internal tool
        ToolFactory.register_tool(MockInternalSearchTool, MockInternalSearchConfig)

        self.event = get_event_obj(
            user_id="test_user",
            company_id="test_company",
            assistant_id="test_assistant",
            chat_id="test_chat",
        )

        # Set tool choices to include both internal and MCP tools
        self.event.payload.tool_choices = ["internal_search", "mcp_test_tool"]
        self.event.payload.disabled_tools = []

    @pytest.fixture
    def mock_chat_service(self):
        """Create mock chat service for tool progress reporter"""
        return Mock(spec=ChatService)

    @pytest.fixture
    def tool_progress_reporter(self, mock_chat_service):
        """Create tool progress reporter fixture"""
        return ToolProgressReporter(mock_chat_service)

    @pytest.fixture
    def mcp_tools(self):
        """Create mock MCP tools fixture"""
        mcp_tool = McpTool(
            id="mcp_test_tool_id",
            name="mcp_test_tool",
            description="Test MCP tool",
            input_schema={
                "$schema": "https://json-schema.org/draft/2020-12/schema",
                "type": "object",
                "properties": {
                    "folder": {
                        "description": "Mail folder to read mails from",
                        "default": "inbox",
                        "type": "string",
                    },
                    "limit": {
                        "description": "Number of emails to retrieve",
                        "default": 10,
                        "type": "number",
                        "minimum": 1,
                        "maximum": 50,
                    },
                },
            },
            output_schema=None,
            annotations=None,
            title="Test MCP Tool",
            icon=None,
            system_prompt=None,
            user_prompt=None,
            is_connected=True,
        )
        return [mcp_tool]

    @pytest.fixture
    def mcp_servers(self, mcp_tools):
        """Create mock MCP servers fixture"""
        server = McpServer(
            id="test_server_id",
            name="test_server",
            description="Test MCP server",
            tools=mcp_tools,
            system_prompt="Test system prompt",
            user_prompt="Test user prompt",
            is_connected=True,
        )
        return [server]

    @pytest.fixture
    def internal_tools(self):
        """Create internal tools fixture"""
        internal_tool_config = ToolBuildConfig(
            name="internal_search",
            configuration=MockInternalSearchConfig(),
            display_name="Internal Search",
            icon=ToolIcon.BOOK,
            selection_policy=ToolSelectionPolicy.BY_USER,
            is_exclusive=False,
            is_enabled=True,
        )
        return [internal_tool_config]

    @pytest.fixture
    def mcp_manager(self, mcp_servers, tool_progress_reporter):
        """Create MCP manager fixture"""
        return MCPManager(
            mcp_servers=mcp_servers,
            event=self.event,
            tool_progress_reporter=tool_progress_reporter,
        )

    @pytest.fixture
    def a2a_manager(self, tool_progress_reporter):
        """Create MCP manager fixture"""
        return A2AManager(
            logger=self.logger,
            tool_progress_reporter=tool_progress_reporter,
            response_watcher=SubAgentResponseWatcher(),
        )

    @pytest.fixture
    def tool_manager_config(self, internal_tools):
        """Create tool manager configuration fixture"""
        return ToolManagerConfig(tools=internal_tools, max_tool_calls=10)

    @pytest.fixture
    def tool_manager(
        self, tool_manager_config, mcp_manager, a2a_manager, tool_progress_reporter
    ):
        """Create tool manager fixture"""

        return ToolManager(
            logger=self.logger,
            config=tool_manager_config,
            event=self.event,
            tool_progress_reporter=tool_progress_reporter,
            mcp_manager=mcp_manager,
            a2a_manager=a2a_manager,
        )

    def test_tool_manager_initialization(self, tool_manager):
        """Test tool manager is initialized correctly"""
        assert tool_manager is not None

        assert (
            len(tool_manager.get_tools()) >= 2
        )  # Should have both internal and MCP tools

    def test_tool_manager_has_both_tool_types(self, tool_manager):
        """Test that tool manager contains both MCP and internal tools"""
        tools = tool_manager.get_tools()
        tool_names = [tool.name for tool in tools]

        # Should contain internal search tool
        assert "internal_search" in tool_names

        # Should contain MCP tool (wrapped)
        assert "mcp_test_tool" in tool_names

        # Should have at least 2 tools total
        assert len(tools) >= 2

    def test_tool_manager_can_get_tools_by_name(self, tool_manager):
        """Test that tool manager can retrieve tools by name"""
        # Test getting internal tool
        internal_tool = tool_manager.get_tool_by_name("internal_search")
        assert internal_tool is not None
        assert internal_tool.name == "internal_search"

        # Test getting MCP tool
        mcp_tool = tool_manager.get_tool_by_name("mcp_test_tool")
        assert mcp_tool is not None
        assert mcp_tool.name == "mcp_test_tool"

    def test_tool_manager_tools_property_contains_both_types(self, tool_manager):
        """Test that the _tools property contains both internal and MCP tools"""
        # Access the private _tools attribute directly to verify integration
        tools = tool_manager._tools
        tool_names = [tool.name for tool in tools]

        # Verify both tool types are present
        assert "internal_search" in tool_names, (
            f"Internal tool missing. Available tools: {tool_names}"
        )
        assert "mcp_test_tool" in tool_names, (
            f"MCP tool missing. Available tools: {tool_names}"
        )

        # Verify we have the expected number of tools
        assert len(tools) == 2, f"Expected 2 tools, got {len(tools)}: {tool_names}"

    def test_tool_manager_logs_loaded_tools(self, tool_manager, caplog):
        """Test that tool manager logs the loaded tools correctly"""
        with caplog.at_level(logging.INFO):
            tool_manager.log_loaded_tools()

        # Check that both tools are mentioned in the logs
        log_output = caplog.text
        assert "internal_search" in log_output
        assert "mcp_test_tool" in log_output

    def test_tool_manager_gets_tool_definitions(self, tool_manager):
        """Test that tool manager returns tool definitions for both tool types"""
        definitions = tool_manager.get_tool_definitions()

        # Should have definitions for both tools
        assert len(definitions) == 2

        definition_names = [defn.name for defn in definitions]
        assert "internal_search" in definition_names
        assert "mcp_test_tool" in definition_names

    def test_init_tools_method(
        self, tool_manager_config, mcp_manager, tool_progress_reporter
    ):
        """Test the _init__tools method behavior with different scenarios"""

        # Test 1: Normal initialization with both tool types

        a2a_manager = A2AManager(
            logger=self.logger,
            tool_progress_reporter=tool_progress_reporter,
            response_watcher=SubAgentResponseWatcher(),
        )

        tool_manager = ToolManager(
            logger=self.logger,
            config=tool_manager_config,
            event=self.event,
            tool_progress_reporter=tool_progress_reporter,
            mcp_manager=mcp_manager,
            a2a_manager=a2a_manager,
        )

        # Verify both tools are loaded
        tools = tool_manager.get_tools()
        tool_names = [tool.name for tool in tools]
        assert "internal_search" in tool_names
        assert "mcp_test_tool" in tool_names
        assert len(tools) == 2

    def test_init_tools_with_disabled_tools(
        self, tool_manager_config, mcp_manager, tool_progress_reporter
    ):
        """Test _init__tools method when some tools are disabled"""

        # Modify event to disable the internal search tool
        event_with_disabled = get_event_obj(
            user_id="test_user",
            company_id="test_company",
            assistant_id="test_assistant",
            chat_id="test_chat",
        )
        event_with_disabled.payload.tool_choices = ["internal_search", "mcp_test_tool"]
        event_with_disabled.payload.disabled_tools = ["internal_search"]

        a2a_manager = A2AManager(
            logger=self.logger,
            tool_progress_reporter=tool_progress_reporter,
            response_watcher=SubAgentResponseWatcher(),
        )

        tool_manager = ToolManager(
            logger=self.logger,
            config=tool_manager_config,
            event=event_with_disabled,
            tool_progress_reporter=tool_progress_reporter,
            mcp_manager=mcp_manager,
            a2a_manager=a2a_manager,
        )

        # Should only have MCP tool, internal tool should be filtered out
        tools = tool_manager.get_tools()
        tool_names = [tool.name for tool in tools]
        assert "internal_search" not in tool_names
        assert "mcp_test_tool" in tool_names
        assert len(tools) == 1

    def test_init_tools_with_limited_tool_choices(
        self, tool_manager_config, mcp_manager, tool_progress_reporter
    ):
        """Test _init__tools method when only specific tools are chosen"""

        # Modify event to only choose internal search tool
        event_with_limited_choices = get_event_obj(
            user_id="test_user",
            company_id="test_company",
            assistant_id="test_assistant",
            chat_id="test_chat",
        )
        event_with_limited_choices.payload.tool_choices = ["internal_search"]
        event_with_limited_choices.payload.disabled_tools = []

        a2a_manager = A2AManager(
            logger=self.logger,
            tool_progress_reporter=tool_progress_reporter,
            response_watcher=SubAgentResponseWatcher(),
        )

        tool_manager = ToolManager(
            logger=self.logger,
            config=tool_manager_config,
            event=event_with_limited_choices,
            tool_progress_reporter=tool_progress_reporter,
            mcp_manager=mcp_manager,
            a2a_manager=a2a_manager,
        )

        # Should only have internal search tool
        tools = tool_manager.get_tools()
        tool_names = [tool.name for tool in tools]
        assert "internal_search" in tool_names
        assert "mcp_test_tool" not in tool_names
        assert len(tools) == 1

    def test_init_tools_with_exclusive_tool(self, mcp_manager, tool_progress_reporter):
        """Test _init__tools method when an exclusive tool is present"""

        # Create an exclusive tool configuration
        exclusive_tool_config = ToolBuildConfig(
            name="internal_search",
            configuration=MockInternalSearchConfig(),
            display_name="Internal Search",
            icon=ToolIcon.BOOK,
            selection_policy=ToolSelectionPolicy.BY_USER,
            is_exclusive=True,  # Make it exclusive
            is_enabled=True,
        )

        config_with_exclusive = ToolManagerConfig(
            tools=[exclusive_tool_config], max_tool_calls=10
        )

        a2a_manager = A2AManager(
            logger=self.logger,
            tool_progress_reporter=tool_progress_reporter,
            response_watcher=SubAgentResponseWatcher(),
        )

        tool_manager = ToolManager(
            logger=self.logger,
            config=config_with_exclusive,
            event=self.event,
            tool_progress_reporter=tool_progress_reporter,
            mcp_manager=mcp_manager,
            a2a_manager=a2a_manager,
        )

        # Should only have the exclusive tool, MCP tools should be ignored
        tools = tool_manager.get_tools()
        tool_names = [tool.name for tool in tools]
        assert "internal_search" in tool_names
        assert "mcp_test_tool" not in tool_names
        assert len(tools) == 1

    def test_init_tools_with_disabled_tool_config(
        self, mcp_manager, tool_progress_reporter
    ):
        """Test _init__tools method when a tool is disabled in its configuration"""

        # Create a disabled tool configuration
        disabled_tool_config = ToolBuildConfig(
            name="internal_search",
            configuration=MockInternalSearchConfig(),
            display_name="Internal Search",
            icon=ToolIcon.BOOK,
            selection_policy=ToolSelectionPolicy.BY_USER,
            is_exclusive=False,
            is_enabled=False,  # Disable the tool
        )

        config_with_disabled = ToolManagerConfig(
            tools=[disabled_tool_config], max_tool_calls=10
        )

        a2a_manager = A2AManager(
            logger=self.logger,
            tool_progress_reporter=tool_progress_reporter,
            response_watcher=SubAgentResponseWatcher(),
        )

        tool_manager = ToolManager(
            logger=self.logger,
            config=config_with_disabled,
            event=self.event,
            tool_progress_reporter=tool_progress_reporter,
            mcp_manager=mcp_manager,
            a2a_manager=a2a_manager,
        )

        # Should only have MCP tool, disabled internal tool should be filtered out
        tools = tool_manager.get_tools()
        tool_names = [tool.name for tool in tools]
        assert "internal_search" not in tool_names
        assert "mcp_test_tool" in tool_names
        assert len(tools) == 1

    def test_mcp_tool_wrapper_generates_correct_parameters(
        self, mcp_tools, mcp_servers, tool_progress_reporter
    ):
        """Test that MCPToolWrapper correctly generates parameters field from MCP tool input schema"""
        # Get the first MCP tool and server from fixtures
        mcp_tool = mcp_tools[0]
        mcp_server = mcp_servers[0]

        # Create MCPToolConfig
        config = MCPToolConfig(
            server_id="test_server_id",
            server_name="test_server",
            mcp_source_id="test_server_id",
        )

        # Initialize the MCPToolWrapper
        tool_wrapper = MCPToolWrapper(
            mcp_server=mcp_server,
            mcp_tool=mcp_tool,
            config=config,
            event=self.event,
            tool_progress_reporter=tool_progress_reporter,
        )

        # Call tool_description to get the LanguageModelToolDescription
        tool_description = tool_wrapper.tool_description()

        # Verify the basic properties
        assert tool_description.name == "mcp_test_tool"
        assert tool_description.description == "Test MCP tool"

        # Verify that parameters field is correctly set to the input_schema
        expected_parameters = {
            "$schema": "https://json-schema.org/draft/2020-12/schema",
            "type": "object",
            "properties": {
                "folder": {
                    "description": "Mail folder to read mails from",
                    "default": "inbox",
                    "type": "string",
                },
                "limit": {
                    "description": "Number of emails to retrieve",
                    "default": 10,
                    "type": "number",
                    "minimum": 1,
                    "maximum": 50,
                },
            },
        }

        assert tool_description.parameters == expected_parameters

        # Verify specific parameter properties
        parameters = tool_description.parameters
        assert isinstance(parameters, dict)
        assert parameters["type"] == "object"
        assert "properties" in parameters

        # Check folder parameter
        folder_prop = parameters["properties"]["folder"]
        assert folder_prop["type"] == "string"
        assert folder_prop["description"] == "Mail folder to read mails from"
        assert folder_prop["default"] == "inbox"

        # Check limit parameter
        limit_prop = parameters["properties"]["limit"]
        assert limit_prop["type"] == "number"
        assert limit_prop["description"] == "Number of emails to retrieve"
        assert limit_prop["default"] == 10
        assert limit_prop["minimum"] == 1
        assert limit_prop["maximum"] == 50
