from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal, overload

from medio.backends.itk_io import ItkIO
from medio.backends.nib_io import NibIO
from medio.backends.pdcm_io import PdcmIO
from medio.metadata.convert_nib_itk import inv_axcodes
from medio.utils.files import is_nifti

if TYPE_CHECKING:
    import os

    import numpy as np
    from numpy.typing import NDArray

    from medio.metadata.metadata import CoordSys, HeaderDict, MetaData

ReadBackend = Literal["itk", "nib", "pdcm", "pydicom"]
WriteBackend = Literal["itk", "nib"]


@overload
def read_img(
    input_path: str | os.PathLike[str],
    desired_ornt: str | None = ...,
    backend: ReadBackend | None = ...,
    dtype: np.dtype[np.generic] | type | None = ...,
    *,
    header: Literal[True],
    channels_axis: int | None = ...,
    coord_sys: CoordSys | None = ...,
    **kwargs: Any,
) -> tuple[NDArray[np.generic], MetaData[HeaderDict]]: ...


@overload
def read_img(
    input_path: str | os.PathLike[str],
    desired_ornt: str | None = ...,
    backend: ReadBackend | None = ...,
    dtype: np.dtype[np.generic] | type | None = ...,
    header: bool = ...,
    channels_axis: int | None = ...,
    coord_sys: CoordSys | None = ...,
    **kwargs: Any,
) -> tuple[NDArray[np.generic], MetaData[object]]: ...


def read_img(
    input_path: str | os.PathLike[str],
    desired_ornt: str | None = None,
    backend: ReadBackend | None = None,
    dtype: np.dtype[np.generic] | type | None = None,
    header: bool = False,
    channels_axis: int | None = -1,
    coord_sys: CoordSys | None = "itk",
    **kwargs: Any,
) -> tuple[NDArray[np.generic], MetaData[object] | MetaData[HeaderDict]]:
    """
    Read medical image with nibabel or itk
    :param input_path: str or os.PathLike, the input path of image file or a directory containing dicom series
    :param desired_ornt: optional parameter for reorienting the image to a desired orientation, e.g. 'RAS'.
    The desired_ornt string is in the convention of `coord_sys` argument (itk by default).
    :param backend: optional parameter for setting the reader backend: 'itk', 'nib', 'pdcm' (also 'pydicom') or None
    :param dtype: equivalent to np_image.astype(dtype) if dtype is not None
    :param header: if True, the returned metadata will include a header attribute with additional metadata dictionary as
    read by the backend. Note: currently, this is supported for files only
    :param channels_axis: if not None and the image is channeled (e.g. RGB) move the channels to channels_axis in the
    returned image array
    :param coord_sys: the coordinate system (or convention) of the `desired_ornt` parameter and the returned metadata.
    It can be 'itk', 'nib' or None, and is 'itk' by default.
    None means that the backend will determine coord_sys, but it can lead to a backend-dependent array and metadata
    :return: numpy image and metadata object
    """
    nib_reader_data = (NibIO.read_img, NibIO.coord_sys)
    itk_reader_data = (ItkIO.read_img, ItkIO.coord_sys)
    pdcm_reader_data = (PdcmIO.read_img, PdcmIO.coord_sys)
    if backend is None:
        if is_nifti(input_path):
            reader, reader_sys = nib_reader_data
        else:
            reader, reader_sys = itk_reader_data
    else:
        if backend == "nib":
            reader, reader_sys = nib_reader_data
        elif backend == "itk":
            reader, reader_sys = itk_reader_data
        elif backend in ("pdcm", "pydicom"):
            reader, reader_sys = pdcm_reader_data
        else:
            raise ValueError('The backend argument must be one of: "itk", "nib", "pdcm" (or "pydicom"), None')

    if (coord_sys is not None) and (coord_sys != reader_sys):
        desired_ornt = inv_axcodes(desired_ornt)

    np_image, metadata = reader(input_path, desired_ornt, header, channels_axis, **kwargs)

    if dtype is not None:
        np_image = np_image.astype(dtype, copy=False)
    if coord_sys is not None:
        metadata.convert(coord_sys)
    return np_image, metadata


def read_meta(
    input_path: str | os.PathLike[str],
    desired_ornt: str | None = None,
    backend: ReadBackend | None = None,
    header: bool = False,
    coord_sys: CoordSys | None = "itk",
    **kwargs: Any,
) -> MetaData[object]:
    """
    Read only the metadata of a medical image (affine, orientation, spatial shape) without loading pixel data.
    :param input_path: str or os.PathLike, the input path of an image file or a DICOM directory
    :param desired_ornt: optional orientation to reorient the returned metadata, e.g. 'RAS'. Uses the `coord_sys`
    convention (itk by default).
    :param backend: optional backend to use: 'itk', 'nib', 'pdcm' (also 'pydicom') or None (auto-detected)
    :param header: if True, the returned metadata will include a header attribute with the raw backend metadata
    dictionary. Currently supported for single files only (not DICOM series).
    :param coord_sys: coordinate system of `desired_ornt` and of the returned metadata: 'itk', 'nib' or None.
    :return: MetaData object with spatial_shape set to the image dimensions
    """
    nib_reader_data = (NibIO.read_meta, NibIO.coord_sys)
    itk_reader_data = (ItkIO.read_meta, ItkIO.coord_sys)
    pdcm_reader_data = (PdcmIO.read_meta, PdcmIO.coord_sys)
    if backend is None:
        if is_nifti(input_path):
            reader_meta, reader_sys = nib_reader_data
        else:
            reader_meta, reader_sys = itk_reader_data
    else:
        if backend == "nib":
            reader_meta, reader_sys = nib_reader_data
        elif backend == "itk":
            reader_meta, reader_sys = itk_reader_data
        elif backend in ("pdcm", "pydicom"):
            reader_meta, reader_sys = pdcm_reader_data
        else:
            raise ValueError('The backend argument must be one of: "itk", "nib", "pdcm" (or "pydicom"), None')

    if (coord_sys is not None) and (coord_sys != reader_sys):
        desired_ornt = inv_axcodes(desired_ornt)

    metadata = reader_meta(input_path, desired_ornt, header, **kwargs)

    if coord_sys is not None:
        metadata.convert(coord_sys)
    return metadata


def save_img(
    filename: str | os.PathLike[str],
    np_image: NDArray[np.generic],
    metadata: MetaData[Any],
    use_original_ornt: bool = True,
    backend: WriteBackend | None = None,
    dtype: np.dtype[np.generic] | type | None = None,
    channels_axis: int | None = None,
    mkdir: bool = False,
    parents: bool = False,
    **kwargs: Any,
) -> None:
    """
    Save numpy image with corresponding metadata to file
    :param filename: str or os.PathLike, the output filename
    :param np_image: the numpy image
    :param metadata: the metadata of the image
    :param use_original_ornt: whether to save in the original orientation stored in metadata.orig_ornt or not
    :param backend: optional - 'itk', 'nib' or None
    :param dtype: equivalent to np_image.astype(dtype) if dtype is not None
    :param channels_axis: if not None - the image is channeled (e.g. RGB) and the channels are in channels_axis
    :param mkdir: if True, creates the directory of `filename`
    :param parents: to be used with `mkdir=True`. If True, creates also the parent directories
    """
    nib_writer = NibIO.save_img
    itk_writer = ItkIO.save_img
    if backend is None:
        writer = nib_writer if is_nifti(filename, check_exist=False) else itk_writer
    else:
        if backend == "nib":
            writer = nib_writer
        elif backend == "itk":
            writer = itk_writer
        else:
            raise ValueError('The backend argument must be one of: "itk", "nib", None')
    if mkdir:
        Path(filename).parent.mkdir(parents=parents, exist_ok=True)
    if dtype is not None:
        np_image = np_image.astype(dtype, copy=False)
    writer(filename, np_image, metadata, use_original_ornt, channels_axis, **kwargs)


def save_dir(
    dirname: str | os.PathLike[str],
    np_image: NDArray[np.generic],
    metadata: MetaData[Any],
    use_original_ornt: bool = True,
    dtype: np.dtype[np.generic] | type | None = None,
    channels_axis: int | None = None,
    parents: bool = False,
    exist_ok: bool = False,
    allow_dcm_reorient: bool = False,
    **kwargs: Any,
) -> None:
    """
    Save image as a dicom directory. See medio.backends.itk_io.ItkIO.save_dcm_dir documentation.
    dtype is equivalent to passing image_np.astype(dtype) if dtype is not None
    """
    if dtype is not None:
        np_image = np_image.astype(dtype, copy=False)
    ItkIO.save_dcm_dir(
        dirname, np_image, metadata, use_original_ornt, channels_axis, parents, exist_ok, allow_dcm_reorient, **kwargs
    )
