"""Tests for ToolRegistry and ToolSelector."""

import pytest

from m4.core.datasets import DatasetDefinition, DatasetRegistry, Modality
from m4.core.tools import Tool, ToolInput, ToolOutput, ToolRegistry, ToolSelector


# Mock tool classes for testing
class MockTabularTool:
    """Mock tool requiring TABULAR modality."""

    name = "mock_tabular"
    description = "Mock tabular tool"
    input_model = ToolInput
    output_model = ToolOutput
    required_modalities = frozenset({Modality.TABULAR})
    supported_datasets = None

    def invoke(self, dataset: DatasetDefinition, params: ToolInput) -> ToolOutput:
        return ToolOutput(result="tabular data")

    def is_compatible(self, dataset: DatasetDefinition) -> bool:
        if self.supported_datasets and dataset.name not in self.supported_datasets:
            return False
        if not self.required_modalities.issubset(dataset.modalities):
            return False
        return True


class MockNotesTool:
    """Mock tool requiring NOTES modality."""

    name = "mock_notes"
    description = "Mock notes tool"
    input_model = ToolInput
    output_model = ToolOutput
    required_modalities = frozenset({Modality.NOTES})
    supported_datasets = None

    def invoke(self, dataset: DatasetDefinition, params: ToolInput) -> ToolOutput:
        return ToolOutput(result="notes data")

    def is_compatible(self, dataset: DatasetDefinition) -> bool:
        if self.supported_datasets and dataset.name not in self.supported_datasets:
            return False
        if not self.required_modalities.issubset(dataset.modalities):
            return False
        return True


class MockMIMICOnlyTool:
    """Mock tool that only works with MIMIC datasets."""

    name = "mock_mimic_only"
    description = "Mock MIMIC-only tool"
    input_model = ToolInput
    output_model = ToolOutput
    required_modalities = frozenset({Modality.TABULAR})
    supported_datasets = frozenset({"mimic-iv", "mimic-iv-demo"})

    def invoke(self, dataset: DatasetDefinition, params: ToolInput) -> ToolOutput:
        return ToolOutput(result="mimic data")

    def is_compatible(self, dataset: DatasetDefinition) -> bool:
        if self.supported_datasets and dataset.name not in self.supported_datasets:
            return False
        if not self.required_modalities.issubset(dataset.modalities):
            return False
        return True


@pytest.fixture(autouse=True)
def reset_registries():
    """Reset tool registry before and after each test."""
    ToolRegistry.reset()
    yield
    ToolRegistry.reset()


class TestToolRegistry:
    """Test ToolRegistry functionality."""

    def test_register_tool(self):
        """Test registering a tool."""
        tool = MockTabularTool()
        ToolRegistry.register(tool)

        registered = ToolRegistry.get("mock_tabular")
        assert registered is not None
        assert registered.name == "mock_tabular"

    def test_register_duplicate_name_raises_error(self):
        """Test that registering a duplicate tool name raises ValueError."""
        tool1 = MockTabularTool()
        tool2 = MockTabularTool()

        ToolRegistry.register(tool1)
        with pytest.raises(ValueError, match="already registered"):
            ToolRegistry.register(tool2)

    def test_get_nonexistent_tool(self):
        """Test getting a tool that doesn't exist."""
        result = ToolRegistry.get("nonexistent")
        assert result is None

    def test_list_all_tools(self):
        """Test listing all registered tools."""
        tool1 = MockTabularTool()
        tool2 = MockNotesTool()

        ToolRegistry.register(tool1)
        ToolRegistry.register(tool2)

        all_tools = ToolRegistry.list_all()
        assert len(all_tools) == 2
        assert tool1 in all_tools
        assert tool2 in all_tools

    def test_list_all_empty(self):
        """Test listing tools when registry is empty."""
        all_tools = ToolRegistry.list_all()
        assert all_tools == []

    def test_reset_clears_registry(self):
        """Test that reset clears all registered tools."""
        ToolRegistry.register(MockTabularTool())
        ToolRegistry.register(MockNotesTool())

        assert len(ToolRegistry.list_all()) == 2

        ToolRegistry.reset()
        assert len(ToolRegistry.list_all()) == 0


class TestToolSelector:
    """Test ToolSelector filtering logic."""

    def test_selector_returns_compatible_tools(self):
        """Test that selector returns only compatible tools."""
        ToolRegistry.register(MockTabularTool())
        ToolRegistry.register(MockMIMICOnlyTool())

        selector = ToolSelector()
        mimic_demo = DatasetRegistry.get("mimic-iv-demo")

        compatible = selector.tools_for_dataset(mimic_demo)

        # Both tools should be compatible with mimic-demo
        assert len(compatible) == 2
        tool_names = {tool.name for tool in compatible}
        assert "mock_tabular" in tool_names
        assert "mock_mimic_only" in tool_names

    def test_selector_filters_by_modality(self):
        """Test selector filters by required modalities."""
        ToolRegistry.register(MockTabularTool())
        ToolRegistry.register(MockNotesTool())

        selector = ToolSelector()
        mimic = DatasetRegistry.get("mimic-iv")

        compatible = selector.tools_for_dataset(mimic)

        # mimic-iv has TABULAR but not NOTES modality
        tool_names = {tool.name for tool in compatible}
        assert "mock_tabular" in tool_names
        assert "mock_notes" not in tool_names

    def test_selector_filters_by_dataset_name(self):
        """Test that selector respects supported_datasets restrictions."""
        ToolRegistry.register(MockMIMICOnlyTool())

        selector = ToolSelector()

        # Should work with MIMIC datasets
        mimic = DatasetRegistry.get("mimic-iv-demo")
        compatible = selector.tools_for_dataset(mimic)
        assert len(compatible) == 1

        # Create a non-MIMIC dataset with same modalities
        eicu = DatasetDefinition(
            name="eicu",
            description="eICU database",
            modalities={Modality.TABULAR},
        )

        # Should NOT work with non-MIMIC datasets (even with modalities)
        compatible = selector.tools_for_dataset(eicu)
        assert len(compatible) == 0

    def test_selector_by_dataset_name_string(self):
        """Test selector using dataset name as string."""
        ToolRegistry.register(MockTabularTool())
        ToolRegistry.register(MockNotesTool())

        selector = ToolSelector()

        # Use string instead of DatasetDefinition
        compatible = selector.tools_for_dataset("mimic-iv")

        # Only tabular tool should match (mimic-iv has TABULAR modality)
        tool_names = {tool.name for tool in compatible}
        assert "mock_tabular" in tool_names

    def test_selector_unknown_dataset_returns_empty(self):
        """Test selector with unknown dataset name."""
        ToolRegistry.register(MockTabularTool())

        selector = ToolSelector()
        compatible = selector.tools_for_dataset("unknown-dataset")

        assert compatible == []

    def test_is_tool_available_by_name(self):
        """Test checking if a specific tool is available."""
        ToolRegistry.register(MockTabularTool())
        ToolRegistry.register(MockMIMICOnlyTool())

        selector = ToolSelector()

        # Tabular tool available for demo (has TABULAR modality)
        assert selector.is_tool_available("mock_tabular", "mimic-iv-demo")

        # MIMIC-only tool is available for demo (it's a MIMIC dataset)
        assert selector.is_tool_available("mock_mimic_only", "mimic-iv-demo")

        # Both available for full
        assert selector.is_tool_available("mock_tabular", "mimic-iv")
        assert selector.is_tool_available("mock_mimic_only", "mimic-iv")

    def test_is_tool_available_with_dataset_definition(self):
        """Test is_tool_available with DatasetDefinition object."""
        ToolRegistry.register(MockTabularTool())

        selector = ToolSelector()
        mimic = DatasetRegistry.get("mimic-iv-demo")

        assert selector.is_tool_available("mock_tabular", mimic)

    def test_is_tool_available_nonexistent_tool(self):
        """Test is_tool_available with tool that doesn't exist."""
        selector = ToolSelector()
        assert not selector.is_tool_available("nonexistent", "mimic-iv-demo")

    def test_is_tool_available_unknown_dataset(self):
        """Test is_tool_available with unknown dataset."""
        ToolRegistry.register(MockTabularTool())

        selector = ToolSelector()
        assert not selector.is_tool_available("mock_tabular", "unknown-dataset")

    def test_tools_for_dataset_empty_registry(self):
        """Test tools_for_dataset when no tools are registered."""
        selector = ToolSelector()
        compatible = selector.tools_for_dataset("mimic-iv-demo")

        assert compatible == []


class TestIntegration:
    """Integration tests for registry and selector together."""

    def test_multiple_tools_with_varying_requirements(self):
        """Test complex scenario with multiple tools and datasets."""
        # Register tools
        ToolRegistry.register(MockTabularTool())
        ToolRegistry.register(MockNotesTool())
        ToolRegistry.register(MockMIMICOnlyTool())

        selector = ToolSelector()

        # Test with demo (has TABULAR modality)
        demo_tools = selector.tools_for_dataset("mimic-iv-demo")
        demo_names = {t.name for t in demo_tools}
        assert "mock_tabular" in demo_names
        assert "mock_mimic_only" in demo_names
        # mock_notes not in demo_names (mimic doesn't have NOTES modality)

        # Test with full (has TABULAR modality)
        full_tools = selector.tools_for_dataset("mimic-iv")
        full_names = {t.name for t in full_tools}
        assert "mock_tabular" in full_names
        assert "mock_mimic_only" in full_names

    def test_tool_protocol_conformance(self):
        """Test that tools conform to the Tool protocol."""
        tool = MockTabularTool()

        # Check protocol conformance
        assert isinstance(tool, Tool)
        assert hasattr(tool, "name")
        assert hasattr(tool, "description")
        assert hasattr(tool, "invoke")
        assert hasattr(tool, "is_compatible")


class TestInitTools:
    """Tests for init_tools function and real tool registration."""

    def test_init_tools_registers_all_tools(self):
        """Test that init_tools registers all expected tools."""
        from m4.core.tools import init_tools, reset_tools

        # Ensure clean state
        reset_tools()

        # Initialize tools
        init_tools()

        # Verify all tools are registered
        all_tools = ToolRegistry.list_all()
        tool_names = {t.name for t in all_tools}

        # Management tools
        assert "list_datasets" in tool_names
        assert "set_dataset" in tool_names

        # Tabular tools
        assert "get_database_schema" in tool_names
        assert "get_table_info" in tool_names
        assert "execute_query" in tool_names

        # Notes tools
        assert "search_notes" in tool_names
        assert "get_note" in tool_names
        assert "list_patient_notes" in tool_names

        # Total: 8 tools (2 management + 3 tabular + 3 notes)
        assert len(all_tools) == 8

        # Cleanup
        reset_tools()

    def test_init_tools_is_idempotent(self):
        """Test that calling init_tools multiple times is safe."""
        from m4.core.tools import init_tools, reset_tools

        reset_tools()

        # Call multiple times
        init_tools()
        init_tools()
        init_tools()

        # Should still have exactly 8 tools
        all_tools = ToolRegistry.list_all()
        assert len(all_tools) == 8

        reset_tools()

    def test_reset_tools_clears_everything(self):
        """Test that reset_tools clears all registered tools."""
        from m4.core.tools import init_tools, reset_tools

        init_tools()
        assert len(ToolRegistry.list_all()) == 8

        reset_tools()
        assert len(ToolRegistry.list_all()) == 0

        # Can reinitialize after reset
        init_tools()
        assert len(ToolRegistry.list_all()) == 8

        reset_tools()

    def test_real_tools_conform_to_protocol(self):
        """Test that all real tool classes conform to the Tool protocol."""
        from m4.core.tools import (
            ExecuteQueryTool,
            GetDatabaseSchemaTool,
            GetNoteTool,
            GetTableInfoTool,
            ListDatasetsTool,
            ListPatientNotesTool,
            SearchNotesTool,
            SetDatasetTool,
        )

        tool_classes = [
            GetDatabaseSchemaTool,
            GetTableInfoTool,
            ExecuteQueryTool,
            ListDatasetsTool,
            SetDatasetTool,
            SearchNotesTool,
            GetNoteTool,
            ListPatientNotesTool,
        ]

        for tool_class in tool_classes:
            tool = tool_class()
            assert isinstance(tool, Tool), f"{tool_class.__name__} is not a Tool"
            assert hasattr(tool, "name")
            assert hasattr(tool, "description")
            assert hasattr(tool, "input_model")
            # output_model removed - tools now return native types
            assert hasattr(tool, "required_modalities")
            assert hasattr(tool, "invoke")
            assert hasattr(tool, "is_compatible")

    def test_selector_with_real_tools(self):
        """Test ToolSelector with the actual registered tools."""
        from m4.core.tools import init_tools, reset_tools

        reset_tools()
        init_tools()

        selector = ToolSelector()

        # Test with demo dataset (has TABULAR modality)
        demo_tools = selector.tools_for_dataset("mimic-iv-demo")
        demo_names = {t.name for t in demo_tools}

        # Management tools should always be available
        assert "list_datasets" in demo_names
        assert "set_dataset" in demo_names

        # Tabular tools should be available for demo
        assert "get_database_schema" in demo_names
        assert "get_table_info" in demo_names
        assert "execute_query" in demo_names

        # Test with full dataset (has TABULAR modality)
        full_tools = selector.tools_for_dataset("mimic-iv")
        full_names = {t.name for t in full_tools}

        # Should have all tools available for full dataset
        assert len(full_names) >= len(demo_names)

        reset_tools()

    def test_management_tools_always_compatible(self):
        """Test that management tools work with any dataset."""
        from m4.core.tools import ListDatasetsTool, SetDatasetTool

        list_tool = ListDatasetsTool()
        set_tool = SetDatasetTool()

        # Create a minimal dataset
        minimal_ds = DatasetDefinition(
            name="minimal",
            modalities=set(),  # No modalities
        )

        # Management tools should be compatible with any dataset
        assert list_tool.is_compatible(minimal_ds)
        assert set_tool.is_compatible(minimal_ds)
