"""DICOM segmentation verification tool."""

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

import pydicom
from mcp.types import ToolAnnotations

from dicom_mcp.server import mcp
from dicom_mcp.constants import ResponseFormat
from dicom_mcp.helpers.tags import _safe_get_tag
from dicom_mcp.helpers.files import _find_dicom_files

# Standard DICOM Segmentation Storage SOP Class UID
_SEG_SOP_CLASS_UID = "1.2.840.10008.5.1.4.1.1.66.4"


def _is_segmentation(ds: pydicom.Dataset) -> bool:
    """Check if a dataset is a segmentation file."""
    sop_class = _safe_get_tag(ds, (0x0008, 0x0016))
    if sop_class == _SEG_SOP_CLASS_UID:
        return True
    if hasattr(ds, "SourceImageSequence"):
        return True
    return False


def _extract_references(ds: pydicom.Dataset) -> List[Dict[str, str]]:
    """Extract referenced SOP Instance UIDs from a segmentation dataset."""
    refs: List[Dict[str, str]] = []

    # Check SourceImageSequence
    if hasattr(ds, "SourceImageSequence"):
        for item in ds.SourceImageSequence:
            ref_uid = _safe_get_tag(item, (0x0008, 0x1155))
            ref_class = _safe_get_tag(item, (0x0008, 0x1150))
            if ref_uid != "N/A":
                refs.append(
                    {
                        "ReferencedSOPInstanceUID": str(ref_uid),
                        "ReferencedSOPClassUID": str(ref_class),
                        "source": "SourceImageSequence",
                    }
                )

    # Check ReferencedSeriesSequence → ReferencedInstanceSequence
    if hasattr(ds, "ReferencedSeriesSequence"):
        for series_item in ds.ReferencedSeriesSequence:
            if hasattr(series_item, "ReferencedInstanceSequence"):
                for inst_item in series_item.ReferencedInstanceSequence:
                    ref_uid = _safe_get_tag(inst_item, (0x0008, 0x1155))
                    ref_class = _safe_get_tag(inst_item, (0x0008, 0x1150))
                    if ref_uid != "N/A":
                        refs.append(
                            {
                                "ReferencedSOPInstanceUID": str(ref_uid),
                                "ReferencedSOPClassUID": str(ref_class),
                                "source": "ReferencedSeriesSequence",
                            }
                        )

    return refs


@mcp.tool(
    name="dicom_verify_segmentations",
    annotations=ToolAnnotations(
        title="Verify Segmentation Source References",
        readOnlyHint=True,
        destructiveHint=False,
        idempotentHint=True,
        openWorldHint=False,
    ),
)
async def dicom_verify_segmentations(
    directory: str,
    recursive: bool = True,
    response_format: ResponseFormat = ResponseFormat.MARKDOWN,
) -> str:
    """Validate that segmentation DICOM files reference valid source images.

    Scans a directory for segmentation files (identified by SOPClassUID or
    SourceImageSequence presence), extracts their referenced SOPInstanceUIDs,
    and verifies that each reference points to an existing source file in
    the same directory.

    Useful for QA of segmentation outputs — ensuring every segmentation
    is linked to a valid source image with no dangling references.

    Args:
        directory: Directory containing segmentation and source DICOM files
        recursive: Search subdirectories (default True)
        response_format: Output format (markdown or json)
    """
    try:
        dir_path = Path(directory)
        if not dir_path.exists():
            return f"Error: Directory not found: {directory}"

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

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

        # Classify files and build source UID index
        seg_files: List[Tuple[Path, pydicom.Dataset]] = []
        source_index: Dict[str, Path] = {}  # SOPInstanceUID → Path

        for fp, ds in all_files:
            sop_uid = _safe_get_tag(ds, (0x0008, 0x0018))
            if sop_uid != "N/A":
                source_index[str(sop_uid)] = fp

            if _is_segmentation(ds):
                seg_files.append((fp, ds))

        if not seg_files:
            return (
                f"No segmentation files found in {directory}. "
                f"Scanned {len(all_files)} DICOM files."
            )

        # Verify references for each segmentation
        results: List[Dict[str, Any]] = []
        total_refs = 0
        matched_refs = 0
        unmatched_refs = 0

        for fp, ds in seg_files:
            seg_sop_uid = _safe_get_tag(ds, (0x0008, 0x0018))
            seg_series_desc = _safe_get_tag(ds, (0x0008, 0x103E))
            references = _extract_references(ds)
            total_refs += len(references)

            ref_results: List[Dict[str, Any]] = []
            for ref in references:
                ref_uid = ref["ReferencedSOPInstanceUID"]
                source_path = source_index.get(ref_uid)
                if source_path:
                    matched_refs += 1
                    ref_results.append(
                        {
                            "ReferencedSOPInstanceUID": ref_uid,
                            "source": ref["source"],
                            "matched": True,
                            "matched_file": source_path.name,
                        }
                    )
                else:
                    unmatched_refs += 1
                    ref_results.append(
                        {
                            "ReferencedSOPInstanceUID": ref_uid,
                            "source": ref["source"],
                            "matched": False,
                            "matched_file": None,
                        }
                    )

            results.append(
                {
                    "segmentation_file": fp.name,
                    "segmentation_sop_uid": str(seg_sop_uid),
                    "series_description": str(seg_series_desc),
                    "reference_count": len(references),
                    "references": ref_results,
                }
            )

        all_matched = unmatched_refs == 0

        if response_format == ResponseFormat.JSON:
            output = {
                "directory": str(dir_path),
                "total_files": len(all_files),
                "truncated": truncated,
                "segmentation_files": len(seg_files),
                "source_files": len(all_files) - len(seg_files),
                "total_references": total_refs,
                "matched_references": matched_refs,
                "unmatched_references": unmatched_refs,
                "all_matched": all_matched,
                "segmentations": results,
            }
            return json.dumps(output, indent=2)
        else:
            status = "ALL MATCHED" if all_matched else "UNMATCHED REFERENCES FOUND"
            output_lines = [
                f"# Segmentation Verification: {dir_path.name}\n",
                f"**Status**: {status}\n",
                "## Summary\n",
                f"- **Total files scanned**: {len(all_files)}",
                f"- **Segmentation files**: {len(seg_files)}",
                f"- **Source files**: {len(all_files) - len(seg_files)}",
                f"- **Total references**: {total_refs}",
                f"- **Matched**: {matched_refs}",
                f"- **Unmatched**: {unmatched_refs}",
            ]

            if truncated:
                output_lines.append(
                    "\n**Warning**: File scan was truncated. "
                    "Results may be incomplete."
                )

            # Show details for segmentations with unmatched references
            unmatched_segs = [
                r for r in results if any(not ref["matched"] for ref in r["references"])
            ]
            if unmatched_segs:
                output_lines.append("\n## Unmatched References\n")
                for seg in unmatched_segs:
                    output_lines.append(f"### {seg['segmentation_file']}")
                    output_lines.append(f"Series: {seg['series_description']}\n")
                    for ref in seg["references"]:
                        if not ref["matched"]:
                            output_lines.append(
                                f"- `{ref['ReferencedSOPInstanceUID']}` "
                                f"(from {ref['source']}) — **NOT FOUND**"
                            )

            if all_matched and seg_files:
                output_lines.append(
                    "\nAll segmentation references successfully matched "
                    "to source files."
                )

            return "\n".join(output_lines)

    except Exception as e:
        return f"Error verifying segmentations: {str(e)}"
