import nibabel as nib
from core.node import Node
from core.contracts_radiomics import ImageDataContract
from pathlib import Path
import uuid


class ImageReader(Node):
    def __init__(self, path: str):
        self.path = Path(path)

        # Read image header only
        img = nib.load(str(self.path))
        header = img.header

        # Infer metadata
        dim = "3D" if img.ndim == 3 else "2D"
        modality = header.get("descrip", b"CT").decode(errors="ignore") or "CT"

        geometry_id = str(uuid.uuid4())

        super().__init__(
            name="ImageReader",
            outputs={
                "image": ImageDataContract(
                    modality=modality,
                    dim=dim,
                    geometry_id=geometry_id,
                    has_mask=False
                )
            }
        )

    def load_data(self):
        # Load full image data only when execution starts
        img = nib.load(str(self.path))
        return img.get_fdata()