"""DICOM tag query and directory summary tools."""

import asyncio
import json
from pathlib import Path
from typing import Any, Dict, List, Optional

from mcp.types import ToolAnnotations

from dicom_mcp.server import mcp
from dicom_mcp.config import MAX_FILES
from dicom_mcp.constants import ResponseFormat
from dicom_mcp.helpers.tags import _safe_get_tag, _resolve_tag, _format_markdown_table
from dicom_mcp.helpers.sequence import _identify_sequence_type
from dicom_mcp.helpers.files import _find_dicom_files
from dicom_mcp.pii import redact_if_pii


@mcp.tool(
    name="dicom_query",
    annotations=ToolAnnotations(
        title="Query DICOM Tags Across Directory",
        readOnlyHint=True,
        destructiveHint=False,
        idempotentHint=True,
        openWorldHint=False,
    ),
)
async def dicom_query(
    directory: str,
    tags: List[str],
    recursive: bool = True,
    group_by: Optional[str] = None,
    response_format: ResponseFormat = ResponseFormat.MARKDOWN,
) -> str:
    """Query arbitrary DICOM tags across all files in a directory.

    Aggregates unique values with file counts for any combination of DICOM
    tags. Accepts tag names (keywords) or hex codes. Use this when you need
    to check consistency or explore values for tags not covered by the
    standard summary tool.

    Args:
        directory: Path to the directory to scan
        tags: List of DICOM tag keywords (e.g. "EchoTime", "StudyDate") or
              hex pairs (e.g. "0018,0081", "0008,0020")
        recursive: Whether to scan subdirectories recursively
        group_by: Optional tag keyword or hex pair to group results by
                  (e.g. "SeriesDescription" to see per-series breakdowns)
        response_format: Output format (markdown or json)
    """
    try:
        dir_path = Path(directory)
        if not dir_path.exists():
            return f"Error: Directory not found: {directory}"
        if not dir_path.is_dir():
            return f"Error: Path is not a directory: {directory}"

        if not tags:
            return "Error: Provide at least one tag to query."

        resolved_tags = []
        errors = []
        for spec in tags:
            tag_tuple, name = _resolve_tag(spec)
            if tag_tuple is None:
                errors.append(f"Could not resolve tag: '{spec}'")
            else:
                resolved_tags.append((tag_tuple, name))

        if errors and not resolved_tags:
            return "Error: " + "; ".join(errors)

        # Resolve group_by tag if provided
        group_by_tag = None
        group_by_name = None
        if group_by:
            group_by_tag, group_by_name = _resolve_tag(group_by)
            if group_by_tag is None:
                return f"Error: Could not resolve group_by tag: '{group_by}'"

        # --- Scan files ---
        dicom_files, truncated = await asyncio.to_thread(
            _find_dicom_files, dir_path, recursive
        )

        if not dicom_files:
            return f"No DICOM files found in {directory}"

        # --- Aggregate values ---
        if group_by_tag:
            # Nested: group_value -> tag_name -> value -> count
            grouped: Dict[str, Dict[str, Dict[str, int]]] = {}
            for _, ds in dicom_files:
                gv = _safe_get_tag(ds, group_by_tag)
                if gv not in grouped:
                    grouped[gv] = {name: {} for _, name in resolved_tags}
                for tag_tuple, name in resolved_tags:
                    value = redact_if_pii(tag_tuple, _safe_get_tag(ds, tag_tuple))
                    grouped[gv][name][value] = grouped[gv][name].get(value, 0) + 1

            if response_format == ResponseFormat.JSON:
                result = {
                    "directory": str(dir_path),
                    "total_files_scanned": len(dicom_files),
                    "truncated": truncated,
                    "group_by": group_by_name,
                    "groups": {},
                }
                if errors:
                    result["warnings"] = errors
                for gv, tag_data in sorted(grouped.items()):
                    result["groups"][gv] = {
                        name: [
                            {"value": v, "count": c}
                            for v, c in sorted(vals.items(), key=lambda x: -x[1])
                        ]
                        for name, vals in tag_data.items()
                    }
                return json.dumps(result, indent=2)
            else:
                output = [
                    f"# DICOM Tag Query: {dir_path}\n",
                    f"**Files scanned**: {len(dicom_files)}",
                    f"**Grouped by**: {group_by_name}\n",
                ]
                if truncated:
                    output.append(
                        f"**Warning**: Scan truncated at {MAX_FILES} files.\n"
                    )
                if errors:
                    output.append(f"**Warnings**: {'; '.join(errors)}\n")
                for gv in sorted(grouped.keys()):
                    output.append(f"## {group_by_name}: {gv}\n")
                    tag_data = grouped[gv]
                    for name, vals in tag_data.items():
                        sorted_vals = sorted(vals.items(), key=lambda x: -x[1])
                        if len(sorted_vals) == 1:
                            v, c = sorted_vals[0]
                            output.append(f"- **{name}**: {v} ({c} files)")
                        else:
                            output.append(f"- **{name}** ({len(sorted_vals)} unique):")
                            for v, c in sorted_vals:
                                output.append(f"  - {v}: {c} files")
                    output.append("")
                return "\n".join(output)
        else:
            # Flat: tag_name -> value -> count
            collectors: Dict[str, Dict[str, int]] = {
                name: {} for _, name in resolved_tags
            }
            for _, ds in dicom_files:
                for tag_tuple, name in resolved_tags:
                    value = redact_if_pii(tag_tuple, _safe_get_tag(ds, tag_tuple))
                    collectors[name][value] = collectors[name].get(value, 0) + 1

            if response_format == ResponseFormat.JSON:
                result = {
                    "directory": str(dir_path),
                    "total_files_scanned": len(dicom_files),
                    "truncated": truncated,
                    "tags": {},
                }
                if errors:
                    result["warnings"] = errors
                for name, vals in collectors.items():
                    result["tags"][name] = {
                        "unique_count": len(vals),
                        "values": [
                            {"value": v, "count": c}
                            for v, c in sorted(vals.items(), key=lambda x: -x[1])
                        ],
                    }
                return json.dumps(result, indent=2)
            else:
                output = [
                    f"# DICOM Tag Query: {dir_path}\n",
                    f"**Files scanned**: {len(dicom_files)}\n",
                ]
                if truncated:
                    output.append(
                        f"**Warning**: Scan truncated at {MAX_FILES} files.\n"
                    )
                if errors:
                    output.append(f"**Warnings**: {'; '.join(errors)}\n")
                for name, vals in collectors.items():
                    sorted_vals = sorted(vals.items(), key=lambda x: -x[1])
                    output.append(f"## {name} ({len(sorted_vals)} unique)\n")
                    headers = ["Value", "File Count"]
                    rows = [[v, str(c)] for v, c in sorted_vals]
                    output.append(_format_markdown_table(headers, rows))
                    output.append("")
                return "\n".join(output)

    except Exception as e:
        return f"Error querying DICOM tags: {str(e)}"


@mcp.tool(
    name="dicom_summarize_directory",
    annotations=ToolAnnotations(
        title="Summarize DICOM Directory",
        readOnlyHint=True,
        destructiveHint=False,
        idempotentHint=True,
        openWorldHint=False,
    ),
)
async def dicom_summarize_directory(
    directory: str,
    recursive: bool = True,
    include_series_overview: bool = True,
    response_format: ResponseFormat = ResponseFormat.MARKDOWN,
) -> str:
    """Summarize a DICOM directory by returning unique values of key tags.

    Scans all DICOM files in a directory and returns distinct values for
    important metadata fields (manufacturer, scanner model, series descriptions,
    sequence types, institution, field strength, etc.) with file counts per
    group. Much more efficient than listing every file when you need an
    overview of what a directory contains.

    Args:
        directory: Path to the directory to scan
        recursive: Whether to scan subdirectories recursively
        include_series_overview: Whether to include detailed series-by-series table
        response_format: Output format (markdown or json)
    """
    try:
        dir_path = Path(directory)
        if not dir_path.exists():
            return f"Error: Directory not found: {directory}"
        if not dir_path.is_dir():
            return f"Error: Path is not a directory: {directory}"

        dicom_files, truncated = await asyncio.to_thread(
            _find_dicom_files, dir_path, recursive
        )

        if not dicom_files:
            return f"No DICOM files found in {directory}"

        # --- Series-level grouping (only if requested) ---
        sorted_series = []
        if include_series_overview:
            series_groups: Dict[str, Dict[str, Any]] = {}
            for _, ds in dicom_files:
                series_uid = _safe_get_tag(ds, (0x0020, 0x000E))
                if series_uid not in series_groups:
                    series_groups[series_uid] = {
                        "series_number": _safe_get_tag(ds, (0x0020, 0x0011)),
                        "series_description": _safe_get_tag(ds, (0x0008, 0x103E)),
                        "sequence_type": _identify_sequence_type(ds).value,
                        "tr": _safe_get_tag(ds, (0x0018, 0x0080)),
                        "te": _safe_get_tag(ds, (0x0018, 0x0081)),
                        "ti": _safe_get_tag(ds, (0x0018, 0x0082)),
                        "fa": _safe_get_tag(ds, (0x0018, 0x1314)),
                        "file_count": 0,
                    }
                series_groups[series_uid]["file_count"] += 1

            # Sort series by series number (numeric where possible)
            def _sort_key(item: tuple) -> tuple:
                series_num = item[1]["series_number"]
                if series_num is None or series_num == "N/A":
                    return (2, "")  # Sort None/N/A values last
                try:
                    return (0, int(series_num))
                except (ValueError, TypeError):
                    return (1, str(series_num))

            sorted_series = sorted(series_groups.items(), key=_sort_key)

        # Tags to summarize: tag tuple -> (display_name, collector dict)
        summary_fields = {
            (0x0010, 0x0010): "Patient Name",
            (0x0010, 0x0020): "Patient ID",
            (0x0010, 0x0030): "Patient Birth Date",
            (0x0010, 0x0040): "Patient Sex",
            (0x0008, 0x0070): "Manufacturer",
            (0x0008, 0x1090): "Scanner Model",
            (0x0018, 0x0087): "Field Strength (T)",
            (0x0008, 0x0080): "Institution",
            (0x0008, 0x1010): "Station Name",
            (0x0008, 0x0060): "Modality",
            (0x0008, 0x103E): "Series Description",
            (0x0020, 0x0011): "Series Number",
        }

        # Collect unique values with counts
        collectors: Dict[str, Dict[str, int]] = {
            name: {} for name in summary_fields.values()
        }
        sequence_type_counts: Dict[str, int] = {}

        for _, ds in dicom_files:
            for tag, field_name in summary_fields.items():
                value = redact_if_pii(tag, _safe_get_tag(ds, tag))
                collectors[field_name][value] = collectors[field_name].get(value, 0) + 1

            seq_type = _identify_sequence_type(ds).value
            sequence_type_counts[seq_type] = sequence_type_counts.get(seq_type, 0) + 1

        # Build summary
        summary = {
            "directory": str(dir_path),
            "total_files_scanned": len(dicom_files),
            "truncated": truncated,
            "fields": {},
        }

        for field_name, value_counts in collectors.items():
            sorted_entries = sorted(value_counts.items(), key=lambda x: -x[1])
            summary["fields"][field_name] = {
                "unique_count": len(sorted_entries),
                "values": [{"value": v, "count": c} for v, c in sorted_entries],
            }

        summary["fields"]["Sequence Type"] = {
            "unique_count": len(sequence_type_counts),
            "values": [
                {"value": v, "count": c}
                for v, c in sorted(sequence_type_counts.items(), key=lambda x: -x[1])
            ],
        }

        # Add series overview to summary (if enabled)
        if include_series_overview and sorted_series:
            summary["series_overview"] = [
                {
                    "series_number": data["series_number"],
                    "series_description": data["series_description"],
                    "sequence_type": data["sequence_type"],
                    "tr": data["tr"],
                    "te": data["te"],
                    "ti": data["ti"],
                    "fa": data["fa"],
                    "file_count": data["file_count"],
                }
                for _, data in sorted_series
            ]

        if response_format == ResponseFormat.JSON:
            return json.dumps(summary, indent=2)
        else:
            output = [
                f"# DICOM Directory Summary: {dir_path}\n",
                f"**Total files scanned**: {len(dicom_files)}",
            ]
            if truncated:
                output.append(
                    f"**Warning**: Scan truncated at {MAX_FILES} files. "
                    "Results may not reflect all data in the directory.\n"
                )
            output.append("")

            # Order fields for readability
            field_order = [
                "Patient Name",
                "Patient ID",
                "Patient Birth Date",
                "Patient Sex",
                "Manufacturer",
                "Scanner Model",
                "Field Strength (T)",
                "Institution",
                "Station Name",
                "Modality",
                "Sequence Type",
                "Series Description",
                "Series Number",
            ]

            for field_name in field_order:
                field_data = summary["fields"].get(field_name)
                if not field_data:
                    continue
                output.append(
                    f"## {field_name} ({field_data['unique_count']} unique)\n"
                )
                headers = ["Value", "File Count"]
                rows = [
                    [entry["value"], str(entry["count"])]
                    for entry in field_data["values"]
                ]
                output.append(_format_markdown_table(headers, rows))
                output.append("")

            # --- Series Overview Table (if enabled) ---
            if include_series_overview and sorted_series:
                output.append("## Series Overview\n")
                series_headers = [
                    "Series",
                    "Description",
                    "Type",
                    "TR",
                    "TE",
                    "TI",
                    "FA",
                    "Files",
                ]
                series_rows = [
                    [
                        data["series_number"],
                        data["series_description"],
                        data["sequence_type"],
                        data["tr"],
                        data["te"],
                        data["ti"],
                        data["fa"],
                        str(data["file_count"]),
                    ]
                    for _, data in sorted_series
                ]
                output.append(_format_markdown_table(series_headers, series_rows))
                output.append("")

            return "\n".join(output)

    except Exception as e:
        return f"Error summarizing directory: {str(e)}"
