import cv2
import os
from workflow.module import Module
from workflow.io_port import OutPort, NIFTIImageType, MaskType
# Codes Go Below:


class ImageReader(Module):
    """
    Pure logical node.
    No runtime, no Dagster, no Engine.
    Only defines ports and run(context) logic.
    """
    def __init__(self):
        super().__init__("ImageReader")

        # Define output ports
        self.addOutPort(OutPort("image", NIFTIImageType()))
        self.addOutPort(OutPort("mask", MaskType()))

    def run(self, context):
        """
        Pure logic: read image + mask from disk.
        No Dagster, no Engine, no state.
        """

        image_path = "data/images/image.nii.gz"
        mask_path = "data/masks/mask.nii.gz"

        if not os.path.exists(image_path):
            raise FileNotFoundError(f"Image not found: {image_path}")

        if not os.path.exists(mask_path):
            raise FileNotFoundError(f"Mask not found: {mask_path}")

        image = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
        mask = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)

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