import os
import numpy as np
import nibabel as nib
import pydicom
import nrrd
import cv2

from api.module import Module
from api.io_port import OutPort, NIFTIImageType


def load_nifti(path):
    nii = nib.load(path)
    return nii.get_fdata().astype(np.float32)


def load_dicom(path):
    if os.path.isdir(path):
        files = sorted([
            os.path.join(path, f)
            for f in os.listdir(path)
            if not f.startswith(".")
        ])
        slices = [pydicom.dcmread(f).pixel_array for f in files]
        return np.stack(slices).astype(np.float32)

    ds = pydicom.dcmread(path)
    return ds.pixel_array.astype(np.float32)


def load_nrrd(path):
    data, _ = nrrd.read(path)
    return data.astype(np.float32)


def load_numpy(path):
    return np.load(path).astype(np.float32)

def load_image(path):
    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
    if img is None:
        raise ValueError(f"Cannot read image file: {path}")
    return img.astype(np.float32)


def load_any(path):
    path = path.replace("\\", "/").lower()

    if path.endswith((".nii", ".nii.gz")):
        return load_nifti(path)

    if path.endswith(".nrrd"):
        return load_nrrd(path)

    if path.endswith(".npy"):
        return load_numpy(path)

    if path.endswith((".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff")):
        return load_image(path)

    if os.path.isdir(path) or path.endswith(".dcm"):
        return load_dicom(path)

    raise ValueError(f"Unsupported image format: {path}")


def factory():
    m = Module("ImageReader")

    out_image = OutPort("image", NIFTIImageType(metadata={"role": "image"}))
    out_mask  = OutPort("mask",  NIFTIImageType(metadata={"role": "mask"}))

    m.add_out_port(out_image)
    m.add_out_port(out_mask)

    def execute(inputs):
        image_dir = "data/images"
        mask_dir  = "data/masks"

        image_files = [f for f in os.listdir(image_dir) if not f.startswith(".")]
        mask_files  = [f for f in os.listdir(mask_dir)  if not f.startswith(".")]

        if not image_files:
            raise ValueError("No image found in data/images")

        if not mask_files:
            raise ValueError("No mask found in data/masks")

        image_path = os.path.join(image_dir, image_files[0])
        mask_path  = os.path.join(mask_dir,  mask_files[0])

        try:
            image = load_any(image_path)
        except Exception as e:
            raise ValueError(f"Cannot read image: {image_path} ({e})")

        try:
            mask = load_any(mask_path)
        except Exception as e:
            raise ValueError(f"Cannot read mask: {mask_path} ({e})")

        return {
            "image": image,
            "mask":  mask
        }

    m.execute = execute
    return m