"""M4 MCP Server - Thin MCP Protocol Adapter.

This module provides the FastMCP server that exposes M4 tools via MCP protocol.
All business logic is delegated to tool classes in m4.core.tools.

Architecture:
    mcp_server.py (this file) - MCP protocol adapter
        ↓ delegates to
    core/tools/*.py - Tool implementations (return native types)
        ↓ uses
    core/backends/*.py - Database backends

    The MCP layer uses the serialization module to convert native Python
    types to MCP-friendly strings. Exceptions are caught and formatted.

Tool Surface:
    The MCP tool surface is stable - all tools remain registered regardless of
    the selected dataset. Compatibility is enforced per-call via proactive
    capability checking before tool invocation.
"""

import json
from typing import Any

import pandas as pd
from fastmcp import FastMCP

# MCP Apps imports
from m4.apps import init_apps
from m4.apps.cohort_builder import RESOURCE_URI as COHORT_BUILDER_URI
from m4.apps.cohort_builder import get_ui_html
from m4.auth import init_oauth2, require_oauth2
from m4.client import M4Client
from m4.core.datasets import DatasetRegistry
from m4.core.exceptions import M4Error
from m4.core.serialization import serialize_for_mcp
from m4.core.telemetry import set_interface
from m4.core.tools import ToolSelector, init_tools
from m4.core.tools.management import ListDatasetsInput, ListDatasetsTool
from m4.services.capabilities import build_capabilities_manifest

# Create FastMCP server instance
mcp = FastMCP("m4")

# Initialize systems
init_oauth2()
init_tools()
init_apps()
set_interface("mcp")

# Tool selector for capability-based filtering
_tool_selector = ToolSelector()

# MCP-exposed tool names (for capability snapshots)
_MCP_TOOL_NAMES = frozenset(
    {
        "list_datasets",
        "get_database_schema",
        "get_table_info",
        "execute_query",
        "search_notes",
        "get_note",
        "list_patient_notes",
        "cohort_builder",
        "query_cohort",
        "capabilities",
    }
)


# ===========================================
# SERIALIZATION HELPERS
# ===========================================


def _serialize_schema_result(result: dict[str, Any]) -> str:
    """Serialize get_database_schema result to MCP string."""
    backend_info = result.get("backend_info", "")
    tables = result.get("tables", [])

    if not tables:
        return f"{backend_info}\n**Available Tables:**\nNo tables found"

    table_list = "\n".join(f"  {t}" for t in tables)
    return f"{backend_info}\n**Available Tables:**\n{table_list}"


def _serialize_table_info_result(result: dict[str, Any]) -> str:
    """Serialize get_table_info result to MCP string."""
    backend_info = result.get("backend_info", "")
    table_name = result.get("table_name", "")
    schema = result.get("schema")
    sample = result.get("sample")

    parts = [
        backend_info,
        f"**Table:** {table_name}",
        "",
        "**Column Information:**",
    ]

    if schema is not None and isinstance(schema, pd.DataFrame):
        parts.append(schema.to_string(index=False))
    else:
        parts.append("(no schema information)")

    if sample is not None and isinstance(sample, pd.DataFrame) and not sample.empty:
        parts.extend(
            [
                "",
                "**Sample Data (first 3 rows):**",
                sample.to_string(index=False),
            ]
        )

    return "\n".join(parts)


def _serialize_datasets_result(result: dict[str, Any]) -> str:
    """Serialize list_datasets result to MCP string."""
    selected = result.get("selected_dataset")
    backend = result.get("backend", "duckdb")
    datasets = result.get("datasets", {})

    if not datasets:
        return "No datasets detected."

    output = []
    if selected:
        output.append(f"Selected dataset: {selected}\n")
    output.append(
        f"Backend: {'local (DuckDB)' if backend == 'duckdb' else 'cloud (BigQuery)'}\n"
    )

    for label, info in datasets.items():
        is_selected = " (Selected)" if info.get("selected") else ""
        output.append(f"=== {label.upper()}{is_selected} ===")

        parquet_icon = "✅" if info.get("parquet_present") else "❌"
        db_icon = "✅" if info.get("db_present") else "❌"
        bq_status = "✅" if info.get("bigquery_support") else "❌"

        output.append(f"  Local Parquet: {parquet_icon}")
        output.append(f"  Local Database: {db_icon}")
        output.append(f"  BigQuery Support: {bq_status}")

        derived = info.get("derived")
        if derived and derived.get("supported"):
            total = derived["total"]
            materialized = derived.get("materialized")
            if materialized is not None:
                icon = "✅" if materialized == total else "⚠️"
                output.append(
                    f"  Derived Tables: {icon} {materialized}/{total} materialized"
                )
                if materialized < total:
                    output.append(f"    Run: m4 init-derived {label}")
            else:
                output.append(f"  Derived Tables: ✅ {total} available")

        output.append("")

    return "\n".join(output)


def _serialize_search_notes_result(result: dict[str, Any]) -> str:
    """Serialize search_notes result to MCP string."""
    backend_info = result.get("backend_info", "")
    query = result.get("query", "")
    snippet_length = result.get("snippet_length", 300)
    results = result.get("results", {})

    if not results or all(df.empty for df in results.values()):
        tables = ", ".join(results.keys()) if results else "notes"
        return f"{backend_info}\n**No matches found** for '{query}' in {tables}."

    output_parts = [
        backend_info,
        f"**Search:** '{query}' (showing snippets of ~{snippet_length} chars)",
    ]

    for table, df in results.items():
        if not df.empty:
            output_parts.append(f"\n**{table.upper()}:**\n{df.to_string(index=False)}")

    output_parts.append(
        "\n**Tip:** Use `get_note(note_id)` to retrieve full text of a specific note."
    )

    return "\n".join(output_parts)


def _serialize_get_note_result(result: dict[str, Any]) -> str:
    """Serialize get_note result to MCP string."""
    backend_info = result.get("backend_info", "")
    note_id = result.get("note_id", "")
    subject_id = result.get("subject_id", "")
    text = result.get("text", "")
    note_length = result.get("note_length", 0)
    truncated = result.get("truncated", False)

    parts = [backend_info, ""]

    if truncated:
        parts.append(f"**Note (truncated, original length: {note_length} chars):**")
    else:
        parts.append(f"**Note {note_id} (subject_id: {subject_id}):**")

    parts.append(text)

    if truncated:
        parts.append("\n[...truncated...]")

    return "\n".join(parts)


def _serialize_list_patient_notes_result(result: dict[str, Any]) -> str:
    """Serialize list_patient_notes result to MCP string."""
    backend_info = result.get("backend_info", "")
    subject_id = result.get("subject_id", "")
    notes = result.get("notes", {})

    if not notes or all(df.empty for df in notes.values()):
        return (
            f"{backend_info}\n**No notes found** for subject_id {subject_id}.\n\n"
            "**Tip:** Verify the subject_id exists in the related MIMIC-IV dataset."
        )

    output_parts = [
        backend_info,
        f"**Notes for subject_id {subject_id}:**",
    ]

    for table, df in notes.items():
        if not df.empty:
            output_parts.append(
                f"\n**{table.upper()} NOTES:**\n{df.to_string(index=False)}"
            )

    output_parts.append(
        "\n**Tip:** Use `get_note(note_id)` to retrieve full text of a specific note."
    )

    return "\n".join(output_parts)


# ==========================================
# MCP TOOLS - Thin adapters to tool classes
# ==========================================


def _client_for_dataset(dataset: str | None) -> M4Client:
    resolved = dataset if dataset is not None else DatasetRegistry.get_active()
    return M4Client(dataset=resolved, interface="mcp")


@mcp.resource("m4://capabilities", mime_type="application/json")
def capabilities_resource() -> str:
    """Return the M4 capability manifest as JSON."""
    return json.dumps(build_capabilities_manifest(), indent=2)


@mcp.tool()
def capabilities() -> str:
    """Return the M4 capability manifest as JSON."""
    return json.dumps(build_capabilities_manifest(), indent=2)


@mcp.tool()
def list_datasets() -> str:
    """📋 List all available datasets and their status.

    Returns:
        A formatted string listing available datasets, indicating selected status,
        and showing availability of local database and BigQuery support.
    """
    try:
        dummy = DatasetRegistry.list_all()[0]
        result = ListDatasetsTool().invoke(dummy, ListDatasetsInput(), context=None)
        return _serialize_datasets_result(result)
    except M4Error as e:
        return f"**Error:** {e}"


@mcp.tool()
def set_dataset(dataset_name: str) -> str:
    """Deprecated: M4 no longer keeps a global active dataset."""
    return (
        "**Error:** set_dataset is no longer supported. M4 tools now require an "
        "explicit dataset argument, for example dataset='mimic-iv'."
    )


@mcp.tool()
@require_oauth2
def get_database_schema(dataset: str | None = None) -> str:
    """📚 Discover what data is available in the database.

    **When to use:** Start here to understand what tables exist.

    Args:
        dataset: Dataset name, e.g. 'mimic-iv'.

    Returns:
        List of all available tables in the database with current backend info.
    """
    try:
        client = _client_for_dataset(dataset)
        dataset_def = client.dataset

        # Proactive capability check
        compat_result = _tool_selector.check_compatibility(
            "get_database_schema", dataset_def
        )
        if not compat_result.compatible:
            return compat_result.error_message

        result = client.schema()
        return _serialize_schema_result(result)
    except M4Error as e:
        return f"**Error:** {e}"


@mcp.tool()
@require_oauth2
def get_table_info(
    table_name: str, show_sample: bool = True, dataset: str | None = None
) -> str:
    """🔍 Explore a specific table's structure and see sample data.

    **When to use:** After identifying relevant tables from get_database_schema().

    Args:
        dataset: Dataset name, e.g. 'mimic-iv'.
        table_name: Exact table name (case-sensitive).
        show_sample: Whether to include sample rows (default: True).

    Returns:
        Table structure with column names, types, and sample data.
    """
    try:
        client = _client_for_dataset(dataset)
        dataset_def = client.dataset

        # Proactive capability check
        compat_result = _tool_selector.check_compatibility(
            "get_table_info", dataset_def
        )
        if not compat_result.compatible:
            return compat_result.error_message

        result = client.table_info(
            table_name=table_name,
            show_sample=show_sample,
        )
        return _serialize_table_info_result(result)
    except M4Error as e:
        return f"**Error:** {e}"


@mcp.tool()
@require_oauth2
def execute_query(sql_query: str, dataset: str | None = None) -> str:
    """🚀 Execute SQL queries to analyze data.

    **Recommended workflow:**
    1. Use get_database_schema() to list tables
    2. Use get_table_info() to examine structure
    3. Write your SQL query with exact names

    Args:
        dataset: Dataset name, e.g. 'mimic-iv'.
        sql_query: Your SQL SELECT query (SELECT only).

    Returns:
        Query results or helpful error messages.
    """
    try:
        client = _client_for_dataset(dataset)
        dataset_def = client.dataset

        # Proactive capability check
        compat_result = _tool_selector.check_compatibility("execute_query", dataset_def)
        if not compat_result.compatible:
            return compat_result.error_message

        result = client.query(sql_query)
        # Result is a DataFrame - serialize it
        return serialize_for_mcp(result)
    except M4Error as e:
        return f"**Error:** {e}"


# ==========================================
# CLINICAL NOTES TOOLS
# ==========================================


@mcp.tool()
@require_oauth2
def search_notes(
    query: str,
    note_type: str = "all",
    limit: int = 5,
    snippet_length: int = 300,
    dataset: str | None = None,
) -> str:
    """🔍 Search clinical notes by keyword.

    Returns snippets around matches to prevent context overflow.
    Use get_note() to retrieve full text of specific notes.

    **Note types:** 'discharge' (summaries), 'radiology' (reports), or 'all'

    Args:
        dataset: Dataset name, e.g. 'mimic-iv-note'.
        query: Search term to find in notes.
        note_type: Type of notes to search ('discharge', 'radiology', or 'all').
        limit: Maximum number of results per note type (default: 5).
        snippet_length: Characters of context around matches (default: 300).

    Returns:
        Matching snippets with note IDs for follow-up retrieval.
    """
    try:
        client = _client_for_dataset(dataset)
        dataset_def = client.dataset

        compat_result = _tool_selector.check_compatibility("search_notes", dataset_def)
        if not compat_result.compatible:
            return compat_result.error_message

        result = client.search_notes(
            query=query,
            note_type=note_type,
            limit=limit,
            snippet_length=snippet_length,
        )
        return _serialize_search_notes_result(result)
    except M4Error as e:
        return f"**Error:** {e}"


@mcp.tool()
@require_oauth2
def get_note(
    note_id: str, max_length: int | None = None, dataset: str | None = None
) -> str:
    """📄 Retrieve full text of a specific clinical note.

    **Warning:** Clinical notes can be very long. Consider using
    search_notes() first to find relevant notes, or use max_length
    to truncate output.

    Args:
        dataset: Dataset name, e.g. 'mimic-iv-note'.
        note_id: The note ID (e.g., from search_notes or list_patient_notes).
        max_length: Optional maximum characters to return (truncates if exceeded).

    Returns:
        Full note text, or truncated version if max_length specified.
    """
    try:
        client = _client_for_dataset(dataset)
        dataset_def = client.dataset

        compat_result = _tool_selector.check_compatibility("get_note", dataset_def)
        if not compat_result.compatible:
            return compat_result.error_message

        result = client.get_note(
            note_id=note_id,
            max_length=max_length,
        )
        return _serialize_get_note_result(result)
    except M4Error as e:
        return f"**Error:** {e}"


@mcp.tool()
@require_oauth2
def list_patient_notes(
    subject_id: int,
    note_type: str = "all",
    limit: int = 20,
    dataset: str | None = None,
) -> str:
    """📋 List available clinical notes for a patient.

    Returns note metadata (IDs, types, lengths) without full text.
    Use get_note(note_id) to retrieve specific notes.

    **Cross-dataset tip:** Get subject_id from MIMIC-IV queries, then
    use it here to find related clinical notes.

    Args:
        dataset: Dataset name, e.g. 'mimic-iv-note'.
        subject_id: Patient identifier (same as in MIMIC-IV).
        note_type: Type of notes to list ('discharge', 'radiology', or 'all').
        limit: Maximum notes to return (default: 20).

    Returns:
        List of available notes with metadata for the patient.
    """
    try:
        client = _client_for_dataset(dataset)
        dataset_def = client.dataset

        compat_result = _tool_selector.check_compatibility(
            "list_patient_notes", dataset_def
        )
        if not compat_result.compatible:
            return compat_result.error_message

        result = client.list_patient_notes(
            subject_id=subject_id,
            note_type=note_type,
            limit=limit,
        )
        return _serialize_list_patient_notes_result(result)
    except M4Error as e:
        return f"**Error:** {e}"


# ==========================================
# MCP APPS - Cohort Builder
# ==========================================


@mcp.resource(COHORT_BUILDER_URI, mime_type="text/html;profile=mcp-app")
def cohort_builder_ui() -> str:
    """Serve the cohort builder UI HTML bundle."""
    return get_ui_html()


@mcp.tool()
@require_oauth2
def cohort_builder(dataset: str | None = None) -> str:
    """Launch the interactive cohort builder.

    Opens a visual interface for filtering patients by demographics and
    clinical criteria. See live patient counts as you adjust filters.

    **Requires:** A host that supports MCP Apps (like Claude Desktop).
    For non-UI hosts, returns dataset information as text.

    Returns:
        Dataset info and welcome message. UI hosts will render the
        interactive cohort builder interface.
    """
    try:
        client = _client_for_dataset(dataset)
        dataset_def = client.dataset

        # Proactive capability check
        compat_result = _tool_selector.check_compatibility(
            "cohort_builder", dataset_def
        )
        if not compat_result.compatible:
            return compat_result.error_message

        result = client.cohort_builder()
        return serialize_for_mcp(result)
    except M4Error as e:
        return f"**Error:** {e}"


@mcp.tool()
@require_oauth2
def query_cohort(
    age_min: int | None = None,
    age_max: int | None = None,
    gender: str | None = None,
    icd_codes: list[str] | None = None,
    icd_match_all: bool | None = None,
    has_icu_stay: bool | None = None,
    in_hospital_mortality: bool | None = None,
    dataset: str | None = None,
) -> str:
    """Query cohort counts based on filtering criteria.

    Used by the cohort builder UI for live updates as users adjust filters.
    Can also be called directly to get cohort statistics.

    Args:
        dataset: Dataset name, e.g. 'mimic-iv'.
        age_min: Minimum patient age (0-130, inclusive).
        age_max: Maximum patient age (0-130, inclusive).
        gender: Patient gender ('M' or 'F').
        icd_codes: List of ICD diagnosis code prefixes to filter by.
        icd_match_all: If True, patient must have ALL ICD codes (AND); default is ANY (OR).
        has_icu_stay: If True, require ICU stay; if False, exclude ICU patients.
        in_hospital_mortality: If True, require in-hospital death; if False, exclude deaths.

    Returns:
        JSON with patient_count, admission_count, demographics, and SQL.
    """
    try:
        client = _client_for_dataset(dataset)
        dataset_def = client.dataset

        # Proactive capability check
        compat_result = _tool_selector.check_compatibility("query_cohort", dataset_def)
        if not compat_result.compatible:
            # Return JSON error for UI compatibility
            return json.dumps({"error": compat_result.error_message})

        result = client.query_cohort(
            age_min=age_min,
            age_max=age_max,
            gender=gender,
            icd_codes=icd_codes,
            icd_match_all=icd_match_all,
            has_icu_stay=has_icu_stay,
            in_hospital_mortality=in_hospital_mortality,
        )
        # Return JSON directly for MCP App UI compatibility
        return json.dumps(result)
    except M4Error as e:
        # Return JSON error for UI compatibility
        return json.dumps({"error": str(e)})


# ==========================================
# _meta.ui.resourceUri INJECTION
# ==========================================


def _inject_cohort_builder_meta() -> None:
    """Inject _meta.ui.resourceUri into the cohort_builder tool.

    FastMCP doesn't expose _meta via the decorator, so we monkey-patch
    the tool's to_mcp_tool method to include it.
    """
    try:
        tool_manager = mcp._tool_manager
        tool_obj = tool_manager._tools.get("cohort_builder")
        if tool_obj is None:
            return

        original_to_mcp = tool_obj.to_mcp_tool

        def patched_to_mcp(**overrides: Any) -> Any:
            overrides.setdefault("_meta", {"ui": {"resourceUri": COHORT_BUILDER_URI}})
            return original_to_mcp(**overrides)

        # Bypass Pydantic's __setattr__ validation
        object.__setattr__(tool_obj, "to_mcp_tool", patched_to_mcp)
    except (AttributeError, TypeError):
        # FastMCP internals may change; fail silently
        pass


# Apply the _meta injection
_inject_cohort_builder_meta()


def main():
    """Main entry point for MCP server."""
    mcp.run()


if __name__ == "__main__":
    main()
