from typing import Set
from core.node import Node
from core.contracts import (
    RequirementContract,
    ImageDataContract,
    FeatureTableContract
)


class ImageRequirement(RequirementContract):
    data_type: str = "image"
    modality: Set[str]
    dim: Set[str]
    requires_mask: bool = False

    def is_satisfied_by(self, provided: ImageDataContract):
        # Validate semantic compatibility
        if provided.data_type != self.data_type:
            raise ValueError("Data type mismatch")

        if provided.modality not in self.modality:
            raise ValueError("Modality mismatch")

        if provided.dim not in self.dim:
            raise ValueError("Dimension mismatch")

        if self.requires_mask and not provided.has_mask:
            raise ValueError("Mask is required but not provided")


class FeatureExtractor(Node):
    def __init__(self):
        super().__init__(
            name="FeatureExtractor",
            inputs={
                "image": ImageRequirement(
                    modality={"CT"},
                    dim={"3D"},
                    requires_mask=False
                )
            },
            outputs={
                "features": FeatureTableContract()
            }
        )
