"""MCP integration tests for cohort builder.

Tests cover:
- cohort_builder tool through fastmcp.Client
- query_cohort tool through fastmcp.Client with test DuckDB
- _meta.ui.resourceUri is present in tool listing
"""

import json
import os
from unittest.mock import patch

import pytest
from fastmcp import Client

from m4.core.datasets import DatasetDefinition, Modality
from m4.core.tools import init_tools
from m4.mcp_server import mcp


def extract_json_from_result(result) -> dict:
    """Extract JSON data from a CallToolResult.

    The result structure has a content attribute with text content.
    """
    # Try to access content attribute
    content = getattr(result, "content", None)
    if content is None:
        # Try direct access
        content = result

    # Find text content
    if isinstance(content, list):
        for item in content:
            if hasattr(item, "text"):
                return json.loads(item.text)
            if isinstance(item, dict) and "text" in item:
                return json.loads(item["text"])

    # Try direct text attribute
    if hasattr(content, "text"):
        return json.loads(content.text)

    raise ValueError(f"Could not extract JSON from result: {result}")


@pytest.fixture(autouse=True)
def ensure_tools_initialized():
    """Ensure tools are initialized before each test."""
    init_tools()


@pytest.fixture
def cohort_test_db(tmp_path):
    """Create a test DuckDB database with cohort-related tables.

    Tables:
    - mimiciv_hosp.patients: subject_id, gender, anchor_age
    - mimiciv_hosp.admissions: subject_id, hadm_id, hospital_expire_flag
    - mimiciv_hosp.diagnoses_icd: hadm_id, icd_code
    - mimiciv_icu.icustays: subject_id, hadm_id, stay_id
    """
    import duckdb

    db_path = tmp_path / "cohort_test.duckdb"
    con = duckdb.connect(str(db_path))
    try:
        # Create schemas
        con.execute("CREATE SCHEMA mimiciv_hosp")
        con.execute("CREATE SCHEMA mimiciv_icu")

        # Create patients table
        con.execute(
            """
            CREATE TABLE mimiciv_hosp.patients (
                subject_id INTEGER PRIMARY KEY,
                gender VARCHAR(1),
                anchor_age INTEGER
            )
            """
        )
        con.execute(
            """
            INSERT INTO mimiciv_hosp.patients (subject_id, gender, anchor_age) VALUES
                (1, 'M', 45),
                (2, 'F', 62),
                (3, 'M', 28),
                (4, 'F', 75),
                (5, 'M', 55)
            """
        )

        # Create admissions table
        con.execute(
            """
            CREATE TABLE mimiciv_hosp.admissions (
                hadm_id INTEGER PRIMARY KEY,
                subject_id INTEGER,
                hospital_expire_flag INTEGER
            )
            """
        )
        con.execute(
            """
            INSERT INTO mimiciv_hosp.admissions (hadm_id, subject_id, hospital_expire_flag) VALUES
                (101, 1, 0),
                (102, 2, 0),
                (103, 3, 1),
                (104, 4, 0),
                (105, 5, 1)
            """
        )

        # Create diagnoses_icd table
        con.execute(
            """
            CREATE TABLE mimiciv_hosp.diagnoses_icd (
                hadm_id INTEGER,
                icd_code VARCHAR(10)
            )
            """
        )
        con.execute(
            """
            INSERT INTO mimiciv_hosp.diagnoses_icd (hadm_id, icd_code) VALUES
                (101, 'I10'),
                (101, 'E11.9'),
                (102, 'I10'),
                (103, 'J18.9'),
                (104, 'E11.9'),
                (105, 'I10')
            """
        )

        # Create icustays table
        con.execute(
            """
            CREATE TABLE mimiciv_icu.icustays (
                stay_id INTEGER PRIMARY KEY,
                subject_id INTEGER,
                hadm_id INTEGER
            )
            """
        )
        con.execute(
            """
            INSERT INTO mimiciv_icu.icustays (stay_id, subject_id, hadm_id) VALUES
                (1001, 1, 101),
                (1002, 3, 103),
                (1003, 5, 105)
            """
        )

        con.commit()
    finally:
        con.close()

    return str(db_path)


class TestCohortBuilderMCPIntegration:
    """Test cohort_builder tool through MCP client."""

    async def test_cohort_builder_tool_exists(self, cohort_test_db):
        """cohort_builder tool should be available."""
        from m4.core.backends import reset_backend_cache

        reset_backend_cache()

        with patch.dict(
            os.environ,
            {
                "M4_BACKEND": "duckdb",
                "M4_DB_PATH": cohort_test_db,
                "M4_OAUTH2_ENABLED": "false",
            },
            clear=True,
        ):
            async with Client(mcp) as client:
                tools = await client.list_tools()
                tool_names = [t.name for t in tools]

                assert "cohort_builder" in tool_names

    async def test_cohort_builder_returns_message(self, cohort_test_db):
        """cohort_builder tool should return welcome message."""
        from m4.core.backends import reset_backend_cache

        reset_backend_cache()

        mock_ds = DatasetDefinition(
            name="mimic-iv-demo",
            modalities=frozenset({Modality.TABULAR}),
        )

        with patch.dict(
            os.environ,
            {
                "M4_BACKEND": "duckdb",
                "M4_DB_PATH": cohort_test_db,
                "M4_OAUTH2_ENABLED": "false",
            },
            clear=True,
        ):
            with patch(
                "m4.mcp_server.DatasetRegistry.get_active", return_value=mock_ds
            ):
                async with Client(mcp) as client:
                    result = await client.call_tool("cohort_builder", {})
                    result_text = str(result)

                    assert (
                        "Cohort Builder" in result_text
                        or "cohort" in result_text.lower()
                    )

    async def test_meta_ui_resource_uri_present(self, cohort_test_db):
        """cohort_builder tool should have _meta.ui.resourceUri."""
        from m4.core.backends import reset_backend_cache

        reset_backend_cache()

        with patch.dict(
            os.environ,
            {
                "M4_BACKEND": "duckdb",
                "M4_DB_PATH": cohort_test_db,
                "M4_OAUTH2_ENABLED": "false",
            },
            clear=True,
        ):
            async with Client(mcp) as client:
                tools = await client.list_tools()

                # Find the cohort_builder tool
                cohort_tool = next(
                    (t for t in tools if t.name == "cohort_builder"), None
                )
                assert cohort_tool is not None

                # Check for _meta.ui.resourceUri
                # The _meta field might be accessed differently depending on MCP version
                meta = getattr(cohort_tool, "meta", None)
                if meta is None:
                    # Try alternative access
                    tool_dict = cohort_tool.model_dump(by_alias=True)
                    meta = tool_dict.get("_meta", {})

                assert meta is not None
                assert "ui" in meta
                assert "resourceUri" in meta["ui"]
                assert "cohort-builder" in meta["ui"]["resourceUri"]


class TestQueryCohortMCPIntegration:
    """Test query_cohort tool through MCP client."""

    async def test_query_cohort_tool_exists(self, cohort_test_db):
        """query_cohort tool should be available."""
        from m4.core.backends import reset_backend_cache

        reset_backend_cache()

        with patch.dict(
            os.environ,
            {
                "M4_BACKEND": "duckdb",
                "M4_DB_PATH": cohort_test_db,
                "M4_OAUTH2_ENABLED": "false",
            },
            clear=True,
        ):
            async with Client(mcp) as client:
                tools = await client.list_tools()
                tool_names = [t.name for t in tools]

                assert "query_cohort" in tool_names

    async def test_query_cohort_empty_criteria(self, cohort_test_db):
        """query_cohort with empty criteria should return all patients."""
        from m4.core.backends import reset_backend_cache

        reset_backend_cache()

        mock_ds = DatasetDefinition(
            name="mimic-iv-demo",
            modalities=frozenset({Modality.TABULAR}),
        )

        with patch.dict(
            os.environ,
            {
                "M4_BACKEND": "duckdb",
                "M4_DB_PATH": cohort_test_db,
                "M4_OAUTH2_ENABLED": "false",
            },
            clear=True,
        ):
            with patch(
                "m4.mcp_server.DatasetRegistry.get_active", return_value=mock_ds
            ):
                with patch("m4.client.get_backend") as mock_backend:
                    from m4.core.backends.duckdb import DuckDBBackend

                    mock_backend.return_value = DuckDBBackend(
                        db_path_override=cohort_test_db
                    )

                    async with Client(mcp) as client:
                        result = await client.call_tool("query_cohort", {})
                        result_text = str(result)

                        # Should have patient count
                        assert "patient_count" in result_text
                        # We have 5 patients in our test DB
                        assert "5" in result_text

    async def test_query_cohort_with_age_filter(self, cohort_test_db):
        """query_cohort with age filter should return filtered patients."""
        from m4.core.backends import reset_backend_cache

        reset_backend_cache()

        mock_ds = DatasetDefinition(
            name="mimic-iv-demo",
            modalities=frozenset({Modality.TABULAR}),
        )

        with patch.dict(
            os.environ,
            {
                "M4_BACKEND": "duckdb",
                "M4_DB_PATH": cohort_test_db,
                "M4_OAUTH2_ENABLED": "false",
            },
            clear=True,
        ):
            with patch(
                "m4.mcp_server.DatasetRegistry.get_active", return_value=mock_ds
            ):
                with patch("m4.client.get_backend") as mock_backend:
                    from m4.core.backends.duckdb import DuckDBBackend

                    mock_backend.return_value = DuckDBBackend(
                        db_path_override=cohort_test_db
                    )

                    async with Client(mcp) as client:
                        # Filter for patients >= 50 years old
                        result = await client.call_tool("query_cohort", {"age_min": 50})

                        # Should return JSON with patient count
                        # Patients >= 50: 2 (62yo), 4 (75yo), 5 (55yo) = 3 patients
                        data = extract_json_from_result(result)
                        assert data["patient_count"] == 3

    async def test_query_cohort_with_gender_filter(self, cohort_test_db):
        """query_cohort with gender filter should return filtered patients."""
        from m4.core.backends import reset_backend_cache

        reset_backend_cache()

        mock_ds = DatasetDefinition(
            name="mimic-iv-demo",
            modalities=frozenset({Modality.TABULAR}),
        )

        with patch.dict(
            os.environ,
            {
                "M4_BACKEND": "duckdb",
                "M4_DB_PATH": cohort_test_db,
                "M4_OAUTH2_ENABLED": "false",
            },
            clear=True,
        ):
            with patch(
                "m4.mcp_server.DatasetRegistry.get_active", return_value=mock_ds
            ):
                with patch("m4.client.get_backend") as mock_backend:
                    from m4.core.backends.duckdb import DuckDBBackend

                    mock_backend.return_value = DuckDBBackend(
                        db_path_override=cohort_test_db
                    )

                    async with Client(mcp) as client:
                        # Filter for female patients
                        result = await client.call_tool("query_cohort", {"gender": "F"})
                        data = extract_json_from_result(result)

                        # Female patients: 2, 4 = 2 patients
                        assert data["patient_count"] == 2

    async def test_query_cohort_with_icd_filter(self, cohort_test_db):
        """query_cohort with ICD code filter should return filtered patients."""
        from m4.core.backends import reset_backend_cache

        reset_backend_cache()

        mock_ds = DatasetDefinition(
            name="mimic-iv-demo",
            modalities=frozenset({Modality.TABULAR}),
        )

        with patch.dict(
            os.environ,
            {
                "M4_BACKEND": "duckdb",
                "M4_DB_PATH": cohort_test_db,
                "M4_OAUTH2_ENABLED": "false",
            },
            clear=True,
        ):
            with patch(
                "m4.mcp_server.DatasetRegistry.get_active", return_value=mock_ds
            ):
                with patch("m4.client.get_backend") as mock_backend:
                    from m4.core.backends.duckdb import DuckDBBackend

                    mock_backend.return_value = DuckDBBackend(
                        db_path_override=cohort_test_db
                    )

                    async with Client(mcp) as client:
                        # Filter for patients with I10 (hypertension)
                        result = await client.call_tool(
                            "query_cohort", {"icd_codes": ["I10"]}
                        )
                        data = extract_json_from_result(result)

                        # Patients with I10: 1, 2, 5 = 3 patients
                        assert data["patient_count"] == 3

    async def test_query_cohort_with_icu_filter(self, cohort_test_db):
        """query_cohort with ICU stay filter should return filtered patients."""
        from m4.core.backends import reset_backend_cache

        reset_backend_cache()

        mock_ds = DatasetDefinition(
            name="mimic-iv-demo",
            modalities=frozenset({Modality.TABULAR}),
        )

        with patch.dict(
            os.environ,
            {
                "M4_BACKEND": "duckdb",
                "M4_DB_PATH": cohort_test_db,
                "M4_OAUTH2_ENABLED": "false",
            },
            clear=True,
        ):
            with patch(
                "m4.mcp_server.DatasetRegistry.get_active", return_value=mock_ds
            ):
                with patch("m4.client.get_backend") as mock_backend:
                    from m4.core.backends.duckdb import DuckDBBackend

                    mock_backend.return_value = DuckDBBackend(
                        db_path_override=cohort_test_db
                    )

                    async with Client(mcp) as client:
                        # Filter for ICU patients only
                        result = await client.call_tool(
                            "query_cohort", {"has_icu_stay": True}
                        )
                        data = extract_json_from_result(result)

                        # ICU patients: 1, 3, 5 = 3 patients
                        assert data["patient_count"] == 3

                        # Filter for non-ICU patients
                        result = await client.call_tool(
                            "query_cohort", {"has_icu_stay": False}
                        )
                        data = extract_json_from_result(result)

                        # Non-ICU patients: 2, 4 = 2 patients
                        assert data["patient_count"] == 2

    async def test_query_cohort_with_mortality_filter(self, cohort_test_db):
        """query_cohort with mortality filter should return filtered patients."""
        from m4.core.backends import reset_backend_cache

        reset_backend_cache()

        mock_ds = DatasetDefinition(
            name="mimic-iv-demo",
            modalities=frozenset({Modality.TABULAR}),
        )

        with patch.dict(
            os.environ,
            {
                "M4_BACKEND": "duckdb",
                "M4_DB_PATH": cohort_test_db,
                "M4_OAUTH2_ENABLED": "false",
            },
            clear=True,
        ):
            with patch(
                "m4.mcp_server.DatasetRegistry.get_active", return_value=mock_ds
            ):
                with patch("m4.client.get_backend") as mock_backend:
                    from m4.core.backends.duckdb import DuckDBBackend

                    mock_backend.return_value = DuckDBBackend(
                        db_path_override=cohort_test_db
                    )

                    async with Client(mcp) as client:
                        # Filter for deceased patients
                        result = await client.call_tool(
                            "query_cohort", {"in_hospital_mortality": True}
                        )
                        data = extract_json_from_result(result)

                        # Deceased patients: 3, 5 = 2 patients
                        assert data["patient_count"] == 2

    async def test_query_cohort_combined_filters(self, cohort_test_db):
        """query_cohort with combined filters should return correct patients."""
        from m4.core.backends import reset_backend_cache

        reset_backend_cache()

        mock_ds = DatasetDefinition(
            name="mimic-iv-demo",
            modalities=frozenset({Modality.TABULAR}),
        )

        with patch.dict(
            os.environ,
            {
                "M4_BACKEND": "duckdb",
                "M4_DB_PATH": cohort_test_db,
                "M4_OAUTH2_ENABLED": "false",
            },
            clear=True,
        ):
            with patch(
                "m4.mcp_server.DatasetRegistry.get_active", return_value=mock_ds
            ):
                with patch("m4.client.get_backend") as mock_backend:
                    from m4.core.backends.duckdb import DuckDBBackend

                    mock_backend.return_value = DuckDBBackend(
                        db_path_override=cohort_test_db
                    )

                    async with Client(mcp) as client:
                        # Male + ICU + I10
                        result = await client.call_tool(
                            "query_cohort",
                            {
                                "gender": "M",
                                "has_icu_stay": True,
                                "icd_codes": ["I10"],
                            },
                        )
                        data = extract_json_from_result(result)

                        # Male ICU patients with I10: 1, 5 = 2 patients
                        assert data["patient_count"] == 2


class TestCohortBuilderResource:
    """Test cohort builder UI resource."""

    async def test_resource_exists(self, cohort_test_db):
        """Cohort builder UI resource should be available."""
        from m4.core.backends import reset_backend_cache

        reset_backend_cache()

        with patch.dict(
            os.environ,
            {
                "M4_BACKEND": "duckdb",
                "M4_DB_PATH": cohort_test_db,
                "M4_OAUTH2_ENABLED": "false",
            },
            clear=True,
        ):
            async with Client(mcp) as client:
                resources = await client.list_resources()
                resource_uris = [r.uri for r in resources]

                # Should have the cohort builder UI resource
                assert any("cohort-builder" in str(uri) for uri in resource_uris)

    async def test_resource_returns_html(self, cohort_test_db):
        """Cohort builder resource should return HTML content."""
        from m4.apps.cohort_builder import get_ui_html

        html = get_ui_html()

        assert "<html" in html.lower()
        assert "M4 Cohort Builder" in html
        # Should have the SDK bundled
        assert "App" in html or "@modelcontextprotocol" in html
