import nibabel as nib
from core.node import Node
from core.compatibility_result import CompatibilityResult, CompatibilityLevel


class ImageWriteRequirement:
    data_type: str = "image"

    def check(self, provided):
        if provided.data_type != self.data_type:
            return CompatibilityResult(
                level=CompatibilityLevel.ERROR,
                message="ImageWriter requires image input"
            )
        return CompatibilityResult(
            level=CompatibilityLevel.OK,
            message="Compatible"
        )


class ImageWriter(Node):
    def __init__(self, output_path: str):
        self.output_path = output_path

        super().__init__(
            name="ImageWriter",
            inputs={"image": ImageWriteRequirement()},
            outputs={}
        )

    def run(self, image_array):
        nib.save(
            nib.Nifti1Image(image_array, None),
            self.output_path
        )