from pydantic import BaseModel
from typing import Literal, Set


class DataContract(BaseModel):
    # Describes what a node produces
    data_type: str
    provides: Set[str] = set()


class ImageDataContract(DataContract):
    data_type: Literal["image"] = "image"
    modality: Literal["CT", "MR"]
    dim: Literal["2D", "3D"]
    has_mask: bool = False
    provides: Set[str] = {"image"}


class MaskDataContract(DataContract):
    data_type: Literal["mask"] = "mask"
    modality: Literal["CT", "MR"]
    dim: Literal["2D", "3D"]
    provides: Set[str] = {"mask"}


class FeatureTableContract(DataContract):
    data_type: Literal["table"] = "table"
    provides: Set[str] = {"features"}


class RequirementContract(BaseModel):
    # Describes what a node requires
    data_type: str

    def is_satisfied_by(self, provided: DataContract):
        raise NotImplementedError
