"""Clinical notes tools for retrieving and searching text data.

This module provides specialized tools for clinical notes that prevent
context overflow by returning snippets instead of full text by default.

Tools:
- search_notes: Full-text search with result snippets
- get_note: Retrieve a single note by ID
- list_patient_notes: List available notes for a patient (metadata only)

Architecture Note:
    Tools return native Python types. The MCP server serializes these
    for the protocol; the Python API receives them directly.
"""

from dataclasses import dataclass
from enum import Enum
from typing import Any

import pandas as pd

from m4.core.backends import get_backend
from m4.core.context import M4ExecutionContext
from m4.core.datasets import DatasetDefinition, Modality
from m4.core.exceptions import QueryError
from m4.core.tools.base import ToolInput


class NoteType(str, Enum):
    """Types of clinical notes available."""

    DISCHARGE = "discharge"
    RADIOLOGY = "radiology"
    ALL = "all"


MAX_NOTE_QUERY_LIMIT = 100
MAX_SNIPPET_LENGTH = 10_000
MAX_NOTE_LENGTH = 200_000
KNOWN_NOTE_TYPES = {item.value for item in NoteType}


def quote_sql_literal(value: str) -> str:
    """Return a single-quoted SQL string literal with quotes escaped."""
    return "'" + value.replace("'", "''") + "'"


def validate_positive_int(value: int, *, name: str, maximum: int) -> int:
    if not isinstance(value, int) or isinstance(value, bool):
        raise QueryError(f"{name} must be an integer.")
    if value <= 0:
        raise QueryError(f"{name} must be positive.")
    if value > maximum:
        raise QueryError(f"{name} must be less than or equal to {maximum}.")
    return value


def validate_note_type(note_type: str) -> str:
    normalized = note_type.lower().strip() if isinstance(note_type, str) else ""
    if normalized not in KNOWN_NOTE_TYPES:
        raise QueryError(
            f"Invalid note_type '{note_type}'. Use 'discharge', 'radiology', or 'all'."
        )
    return normalized


def validate_non_empty_text(value: str, *, name: str) -> str:
    if not isinstance(value, str) or not value.strip():
        raise QueryError(f"{name} must be a non-empty string.")
    return value.strip()


# Input models
@dataclass
class SearchNotesInput(ToolInput):
    """Input for search_notes tool."""

    query: str
    note_type: str = "all"  # discharge, radiology, or all
    limit: int = 5
    snippet_length: int = 300


@dataclass
class GetNoteInput(ToolInput):
    """Input for get_note tool."""

    note_id: str
    max_length: int | None = None  # Optional truncation


@dataclass
class ListPatientNotesInput(ToolInput):
    """Input for list_patient_notes tool."""

    subject_id: int
    note_type: str = "all"  # discharge, radiology, or all
    limit: int = 20


# Tool implementations
class SearchNotesTool:
    """Tool for full-text search across clinical notes.

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

    Returns:
        dict with search results by table, each containing a DataFrame
    """

    name = "search_notes"
    description = (
        "Search clinical notes by keyword. Returns snippets, not full text. "
        "Use get_note() to retrieve full content of a specific note."
    )
    input_model = SearchNotesInput

    required_modalities: frozenset[Modality] = frozenset({Modality.NOTES})
    supported_datasets: frozenset[str] | None = None

    def invoke(
        self,
        dataset: DatasetDefinition,
        params: SearchNotesInput,
        context: M4ExecutionContext | None = None,
    ) -> dict[str, Any]:
        """Search notes and return snippets around matches.

        Returns:
            dict with:
                - backend_info: str - Backend description
                - query: str - Search term used
                - snippet_length: int - Snippet length
                - results: dict[str, pd.DataFrame] - Results by table type

        Raises:
            QueryError: If note_type is invalid
        """
        backend = context.backend if context else get_backend()
        query = validate_non_empty_text(params.query, name="query")
        note_type = validate_note_type(params.note_type)
        limit = validate_positive_int(
            params.limit, name="limit", maximum=MAX_NOTE_QUERY_LIMIT
        )
        snippet_length = validate_positive_int(
            params.snippet_length,
            name="snippet_length",
            maximum=MAX_SNIPPET_LENGTH,
        )

        # Determine which tables to search
        tables_to_search = self._get_tables_for_type(note_type)

        results: dict[str, pd.DataFrame] = {}
        errors: list[str] = []
        search_literal = quote_sql_literal(query)
        search_like_literal = quote_sql_literal(f"%{query.lower()}%")

        for table in tables_to_search:
            # Build search query with snippet extraction
            # Using LIKE for basic search - could be enhanced with full-text search
            sql = f"""
                SELECT
                    note_id,
                    subject_id,
                    CASE
                        WHEN STRPOS(LOWER(text), LOWER({search_literal})) > 0 THEN
                            SUBSTRING(
                                text,
                                GREATEST(1, STRPOS(LOWER(text), LOWER({search_literal})) - {snippet_length // 2}),
                                {snippet_length}
                            )
                        ELSE LEFT(text, {snippet_length})
                    END as snippet,
                    LENGTH(text) as note_length
                FROM {table}
                WHERE LOWER(text) LIKE {search_like_literal}
                LIMIT {limit}
            """

            result = (
                backend.execute_query(sql, dataset, context)
                if context
                else backend.execute_query(sql, dataset)
            )
            if result.success and result.dataframe is not None:
                results[table] = result.dataframe
            elif result.error:
                errors.append(f"{table}: {result.error}")

        response: dict[str, Any] = {
            "backend_info": (
                backend.get_backend_info(dataset, context)
                if context
                else backend.get_backend_info(dataset)
            ),
            "query": query,
            "snippet_length": snippet_length,
            "results": results,
        }
        if errors:
            response["errors"] = errors
        return response

    def _get_tables_for_type(self, note_type: str) -> list[str]:
        """Get table names for a note type."""
        note_type = note_type.lower()
        if note_type == "discharge":
            return ["mimiciv_note.discharge"]
        elif note_type == "radiology":
            return ["mimiciv_note.radiology"]
        elif note_type == "all":
            return ["mimiciv_note.discharge", "mimiciv_note.radiology"]
        return []

    def is_compatible(self, dataset: DatasetDefinition) -> bool:
        """Check compatibility."""
        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 GetNoteTool:
    """Tool for retrieving a single clinical note by ID.

    Returns the full note text. Use with caution as notes can be long.

    Returns:
        dict with note metadata and full text
    """

    name = "get_note"
    description = (
        "Retrieve full text of a specific clinical note by note_id. "
        "Notes can be very long - consider using search_notes() first "
        "to find relevant notes."
    )
    input_model = GetNoteInput

    required_modalities: frozenset[Modality] = frozenset({Modality.NOTES})
    supported_datasets: frozenset[str] | None = None

    def invoke(
        self,
        dataset: DatasetDefinition,
        params: GetNoteInput,
        context: M4ExecutionContext | None = None,
    ) -> dict[str, Any]:
        """Retrieve a single note by ID.

        Returns:
            dict with:
                - backend_info: str - Backend description
                - note_id: str - Note identifier
                - subject_id: int - Patient ID
                - text: str - Full note text (possibly truncated)
                - note_length: int - Original note length
                - truncated: bool - Whether text was truncated

        Raises:
            QueryError: If note not found
        """
        backend = context.backend if context else get_backend()
        note_id = validate_non_empty_text(params.note_id, name="note_id")
        max_length = params.max_length
        if max_length is not None:
            max_length = validate_positive_int(
                max_length,
                name="max_length",
                maximum=MAX_NOTE_LENGTH,
            )

        note_id_literal = quote_sql_literal(note_id)

        # Try both tables since we may not know which one contains the note
        errors: list[str] = []
        for table in ["mimiciv_note.discharge", "mimiciv_note.radiology"]:
            sql = f"""
                SELECT
                    note_id,
                    subject_id,
                    text,
                    LENGTH(text) as note_length
                FROM {table}
                WHERE note_id = {note_id_literal}
                LIMIT 1
            """

            result = (
                backend.execute_query(sql, dataset, context)
                if context
                else backend.execute_query(sql, dataset)
            )
            if result.error:
                errors.append(f"{table}: {result.error}")
                continue
            if (
                result.success
                and result.dataframe is not None
                and not result.dataframe.empty
            ):
                row = result.dataframe.iloc[0]
                text = str(row["text"])
                note_length = int(row["note_length"])
                truncated = False

                # Optionally truncate if max_length specified
                if max_length and len(text) > max_length:
                    text = text[:max_length]
                    truncated = True

                return {
                    "backend_info": (
                        backend.get_backend_info(dataset, context)
                        if context
                        else backend.get_backend_info(dataset)
                    ),
                    "note_id": str(row["note_id"]),
                    "subject_id": int(row["subject_id"]),
                    "text": text,
                    "note_length": note_length,
                    "truncated": truncated,
                }

        error_msg = f"Note '{note_id}' not found."
        if errors:
            error_msg += " Backend errors: " + "; ".join(errors)
        error_msg += (
            " Use list_patient_notes(subject_id) or search_notes(query) "
            "to find valid note IDs."
        )
        raise QueryError(error_msg)

    def is_compatible(self, dataset: DatasetDefinition) -> bool:
        """Check compatibility."""
        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 ListPatientNotesTool:
    """Tool for listing available notes for a patient.

    Returns metadata only (note IDs, types, lengths) - not full text.
    Use this to discover what notes exist before retrieving them.

    Returns:
        dict with notes metadata by table type
    """

    name = "list_patient_notes"
    description = (
        "List available clinical notes for a patient by subject_id. "
        "Returns note metadata (IDs, types, lengths) without full text. "
        "Use get_note(note_id) to retrieve specific notes."
    )
    input_model = ListPatientNotesInput

    required_modalities: frozenset[Modality] = frozenset({Modality.NOTES})
    supported_datasets: frozenset[str] | None = None

    def invoke(
        self,
        dataset: DatasetDefinition,
        params: ListPatientNotesInput,
        context: M4ExecutionContext | None = None,
    ) -> dict[str, Any]:
        """List notes for a patient without returning full text.

        Returns:
            dict with:
                - backend_info: str - Backend description
                - subject_id: int - Patient ID
                - notes: dict[str, pd.DataFrame] - Notes metadata by table type

        Raises:
            QueryError: If note_type is invalid
        """
        backend = context.backend if context else get_backend()
        note_type = validate_note_type(params.note_type)
        limit = validate_positive_int(
            params.limit, name="limit", maximum=MAX_NOTE_QUERY_LIMIT
        )
        subject_id = validate_positive_int(
            params.subject_id, name="subject_id", maximum=9_999_999_999
        )

        tables_to_query = self._get_tables_for_type(note_type)

        notes: dict[str, pd.DataFrame] = {}
        errors: list[str] = []

        for table in tables_to_query:
            # Query for metadata only - explicitly exclude full text
            sql = f"""
                SELECT
                    note_id,
                    subject_id,
                    '{table}' as note_type,
                    LENGTH(text) as note_length,
                    LEFT(text, 100) as preview
                FROM {table}
                WHERE subject_id = {subject_id}
                LIMIT {limit}
            """

            result = (
                backend.execute_query(sql, dataset, context)
                if context
                else backend.execute_query(sql, dataset)
            )
            if result.success and result.dataframe is not None:
                notes[table] = result.dataframe
            elif result.error:
                errors.append(f"{table}: {result.error}")

        response: dict[str, Any] = {
            "backend_info": (
                backend.get_backend_info(dataset, context)
                if context
                else backend.get_backend_info(dataset)
            ),
            "subject_id": subject_id,
            "notes": notes,
        }
        if errors:
            response["errors"] = errors
        return response

    def _get_tables_for_type(self, note_type: str) -> list[str]:
        """Get table names for a note type."""
        note_type = note_type.lower()
        if note_type == "discharge":
            return ["mimiciv_note.discharge"]
        elif note_type == "radiology":
            return ["mimiciv_note.radiology"]
        elif note_type == "all":
            return ["mimiciv_note.discharge", "mimiciv_note.radiology"]
        return []

    def is_compatible(self, dataset: DatasetDefinition) -> bool:
        """Check compatibility."""
        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
