from __future__ import annotations

from collections import Counter
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Dict, List, Optional

from pydicom import dcmread
from sqlalchemy import JSON, Column, ForeignKey, String, Table
from sqlalchemy.orm import relationship
from sqlalchemy.orm.decl_api import registry

from imgtools.logging import logger

if TYPE_CHECKING:
    from pathlib import Path

mapper_registry = registry()


def repr_mixin(cls) -> type:  # type: ignore  # noqa: ANN001
    """
    A decorator to add a customizable __repr__ method to SQLAlchemy ORM models.

    Includes scalar fields (e.g., columns) and prints the length of lists for relationships.
    """

    def __repr_method__(self) -> str:  # type: ignore  # noqa: N807, ANN001
        # Fetch attributes and handle lists separately
        attrs = {
            key: (len(value) if isinstance(value, list) else value)
            for key, value in vars(self).items()
            if not key.startswith("_") and isinstance(value, (str, int, float, list))
        }
        class_name = self.__class__.__name__
        fields = ",\n  ".join(f"{key}={value}" for key, value in attrs.items())
        return f"<{class_name}(\n  {fields}\n)>"

    cls.__repr__ = __repr_method__
    return cls


####################################################################################################
# PATIENT
####################################################################################################


# Patient Table
@repr_mixin
@mapper_registry.mapped
@dataclass
class Patient:
    __table__ = Table(
        "patients",
        mapper_registry.metadata,
        Column("PatientID", String, primary_key=True, index=True),  # PRIMARY KEY
    )

    # Metadata from DICOM
    PatientID: str = field(init=True)

    # Relationships
    studies: List[Study] = field(default_factory=list)

    __mapper_args__ = {  # type: ignore
        "properties": {
            "studies": relationship("Study", back_populates="patient"),
        }
    }

    @classmethod
    def from_metadata(cls, metadata: Dict[str, str]) -> Patient:
        """
        Create a Patient instance from metadata.

        Parameters
        ----------
        metadata : Dict[str, Optional[str]]
                        Patient metadata.

        Returns
        -------
        Patient
                        A Patient instance.
        """
        return cls(
            PatientID=metadata["PatientID"],
        )

    @property
    def num_studies(self) -> int:
        """
        Get the number of studies for the patient.

        Returns
        -------
        int
                        The number of studies.
        """
        return len(self.studies)

    @property
    def num_series(self) -> int:
        """
        Get the number of series for the patient.

        Returns
        -------
        int
                        The number of series.
        """
        return sum(len(study.series) for study in self.studies)

    @property
    def modalities(self) -> Dict[str, int]:
        """
        Get the list of modalities for the patient.

        Returns
        -------
        List[str]
                        The list of modalities.
        """
        modalities = [modality for study in self.studies for modality in study.modalities]
        return dict(Counter(modalities))

    @property
    def series(self) -> List[Series]:
        """
        Get the series for the patient.

        Returns
        -------
        List[Series]
                        The series.
        """
        return [series for study in self.studies for series in study.series]

    @property
    def rtstructs(self) -> List[Series]:
        """
        Get the RTSTRUCT series for the patient.

        Returns
        -------
        List[Series]
                        The RTSTRUCT series.
        """
        return [series for series in self.series if series.Modality == "RTSTRUCT"]


####################################################################################################
# STUDY
####################################################################################################


# Study Table
@repr_mixin
@mapper_registry.mapped
@dataclass
class Study:
    __table__ = Table(
        "studies",
        mapper_registry.metadata,
        Column("StudyInstanceUID", String, primary_key=True, index=True),  # PRIMARY KEY
        Column("PatientID", String, ForeignKey("patients.PatientID"), nullable=False),
    )

    # Metadata from DICOM
    StudyInstanceUID: str = field(init=True)  # PRIMARY KEY
    PatientID: str = field(init=True)  # FOREIGN KEY

    # Relationships
    patient: Patient = field(init=False)
    series: List[Series] = field(default_factory=list)

    __mapper_args__ = {  # type: ignore
        "properties": {
            "patient": relationship("Patient", back_populates="studies"),
            "series": relationship("Series", back_populates="study"),
        }
    }

    @classmethod
    def from_metadata(cls, metadata: Dict[str, str]) -> Study:
        """
        Create a Study instance from metadata.

        Parameters
        ----------
        metadata : Dict[str, Optional[str]]
                        Study metadata.

        Returns
        -------
        Study
                        A Study instance.
        """
        return cls(
            StudyInstanceUID=metadata["StudyInstanceUID"],
            PatientID=metadata["PatientID"],
        )

    @property
    def num_files(self) -> int:
        """
        Get the number of files in the study.

        Returns
        -------
        int
                        The number of files.
        """
        return sum(series.num_files for series in self.series)

    @property
    def modalities(self) -> Dict[str, int]:
        """
        Get the list of modalities in the study.

        Returns
        -------
        List[str]
                        The list of modalities.
        """

        modalities = [series.Modality for series in self.series]
        return dict(Counter(modalities))

    @property
    def unique_modalities(self) -> List[str]:
        """
        Get the unique modalities in the study.

        Returns
        -------
        List[str]
                        The unique modalities.
        """
        return list(self.modalities.keys())


####################################################################################################
# SERIES, RTSTRUCTSERIES, ...
####################################################################################################


@repr_mixin
@mapper_registry.mapped
@dataclass
class Series:
    __table__ = Table(
        "series",
        mapper_registry.metadata,
        Column("SeriesInstanceUID", String, primary_key=True, index=True),  # PRIMARY KEY
        Column("StudyInstanceUID", String, ForeignKey("studies.StudyInstanceUID"), nullable=False),
        Column("Modality", String, nullable=False),
        Column("_RTSTRUCT", JSON, nullable=True),
        Column("_MR", JSON, nullable=True),
    )

    # Metadata from DICOM
    SeriesInstanceUID: str = field(init=True)  # PRIMARY KEY
    StudyInstanceUID: str = field(init=True)  # FOREIGN KEY
    Modality: str = field(init=True)
    _RTSTRUCT: Optional[Dict[str, str]] = field(default_factory=dict)
    _MR: Optional[Dict[str, str]] = field(default_factory=dict)

    # Relationships
    study: Study = field(init=False)
    images: List[Image] = field(default_factory=list)

    __mapper_args__ = {  # type: ignore
        "properties": {
            "study": relationship("Study", back_populates="series"),
            "images": relationship("Image", back_populates="series"),
        },
    }

    @classmethod
    def from_metadata(cls, metadata: Dict[str, str], file_path: Path) -> Series:
        base_cls = cls(
            SeriesInstanceUID=metadata["SeriesInstanceUID"],
            StudyInstanceUID=metadata["StudyInstanceUID"],
            Modality=metadata["Modality"],
        )

        match metadata["Modality"]:
            case "RTSTRUCT":
                try:
                    ds = dcmread(file_path, stop_before_pixels=True)
                    logger.info(f"Reading RTSTRUCT file: {file_path}")
                    roi_names = [roi.ROIName for roi in ds.StructureSetROISequence]
                    rfrs = ds.ReferencedFrameOfReferenceSequence[0]
                    ref_series_seq = (
                        rfrs.RTReferencedStudySequence[0]
                        .RTReferencedSeriesSequence[0]
                        .SeriesInstanceUID
                    )
                    frame_of_reference = rfrs.FrameOfReferenceUID
                    base_cls._RTSTRUCT = {
                        "ROINames": ",".join(roi_names),
                        "RTReferencedSeriesUID": ref_series_seq,
                        "FrameOfReferenceUID": frame_of_reference,
                    }
                except Exception as e:
                    logger.exception(f"Error reading RTSTRUCT file: {file_path}", error=str(e))
                    raise e
            case "MR":
                try:
                    ds = dcmread(file_path, stop_before_pixels=True)
                    logger.info(f"Reading MR file: {file_path}")
                    _mr = {
                        "RepetitionTime": ds.RepetitionTime,
                        "EchoTime": ds.EchoTime,
                        "SliceThickness": ds.SliceThickness,
                        "ScanningSequence": ds.ScanningSequence,
                        "MagneticFieldStrength": ds.MagneticFieldStrength,
                        "ImagedNucleus": ds.ImagedNucleus,
                    }

                    base_cls._MR = {k: str(v) for k, v in _mr.items()}
                except Exception as e:
                    logger.exception(f"Error reading MR file: {file_path}", error=str(e))
                    raise e
            case _:
                pass

        return base_cls

    @property
    def num_files(self) -> int:
        return len(self.images)

    @property
    def RTReferencedSeriesUID(self) -> Optional[str]:  # noqa: N802
        return self._RTSTRUCT.get("RTReferencedSeriesUID", None) if self._RTSTRUCT else None

    @property
    def ROINames(self) -> Optional[str]:  # noqa: N802
        return self._RTSTRUCT.get("ROINames", None) if self._RTSTRUCT else None


####################################################################################################
# IMAGE
####################################################################################################


# Image Table
@repr_mixin
@mapper_registry.mapped
@dataclass
class Image:
    __table__ = Table(
        "images",
        mapper_registry.metadata,
        Column("FilePath", String, primary_key=True, index=True),  # PRIMARY KEY
        Column("SOPInstanceUID", String, nullable=False),
        Column("SeriesInstanceUID", String, ForeignKey("series.SeriesInstanceUID"), nullable=False),
    )

    # Metadata from DICOM
    FilePath: str = field(init=True)
    SOPInstanceUID: str = field(init=True)
    SeriesInstanceUID: str = field(init=True)

    # Relationships
    series: Series = field(init=False)

    __mapper_args__ = {  # type: ignore
        "properties": {
            "series": relationship("Series", back_populates="images"),
        }
    }

    @classmethod
    def from_metadata(cls, metadata: Dict[str, str], file_path: Path) -> Image:
        """
        Create a Image instance from metadata.

        Parameters
        ----------
        metadata : Dict[str, Optional[str]]
                        Image metadata.

        Returns
        -------
        Image
                        A Image instance.
        """
        return cls(
            SOPInstanceUID=metadata["SOPInstanceUID"],
            SeriesInstanceUID=metadata["SeriesInstanceUID"],
            FilePath=file_path.as_posix(),
        )
