from radiuma_api.workflow import Workflow
from radiuma_api.execution_context import ExecutionContext
from radiuma_api.io_port import InPort, OutPort, ImageType, TextType
from radiuma_api.module import Module
from radiuma_api.task import Status
from radiuma_api.asset import Asset
# Codes Go Below:


class Reader(Module):
    def __init__(self):
        super().__init__("Reader")
        self._add_output_port(OutPort("images", ImageType()))

    def run(self, context):
        print("[Reader] Loading images...")
        context.put(self.outputs[0], Asset(["img1", "img2"], "images"))
        self._set_status(Status.COMPLETED)


class Extractor(Module):
    def __init__(self):
        super().__init__("Extractor")
        self._add_input_port(InPort("images", ImageType()))
        self._add_output_port(OutPort("features", TextType()))

    def run(self, context):
        imgs = context.get(self.inputs[0]).data
        print(f"[Extractor] Extracting features from: {imgs}")
        feats = [{"img": img, "feat": f"feature_of_{img}"} for img in imgs]
        context.put(self.outputs[0], Asset(feats, "features"))
        self._set_status(Status.COMPLETED)


class Writer(Module):
    def __init__(self):
        super().__init__("Writer")
        self._add_input_port(InPort("features", TextType()))

    def run(self, context):
        feats = context.get(self.inputs[0]).data
        print("[Writer] Saving features:")
        for f in feats:
            print("   ", f)
        self._set_status(Status.COMPLETED)


def main():
    ctx = ExecutionContext()
    wf = Workflow("RND_Workflow")

    r = Reader()
    e = Extractor()
    w = Writer()

    r.outputs[0].connect(e.inputs[0])
    e.outputs[0].connect(w.inputs[0])

    wf.add_child(r)
    wf.add_child(e)
    wf.add_child(w)

    print("=== Running Workflow ===")
    wf.run(ctx)
    print("=== Workflow Completed ===")


if __name__ == "__main__":
    main()
