from typing_extensions import override
from comfy_api.latest import ComfyExtension, io


class ImagePicker(io.ComfyNode):

    @classmethod
    def define_schema(cls) -> io.Schema:
        return io.Schema(
            node_id="ModelChiefImagePicker",
            display_name="Image Picker",
            category="Model Chief",
            inputs=[
                io.Image.Input("image_1"),
                io.Image.Input("image_2"),
                io.Image.Input("image_3"),
                io.Combo.Input(
                    "pick",
                    options=["Image 1", "Image 2", "Image 3"],
                ),
            ],
            outputs=[
                io.Image.Output(),
            ],
        )

    @classmethod
    def execute(
        cls,
        image_1,
        image_2,
        image_3,
        pick,
    ) -> io.NodeOutput:

        images = {
            "Image 1": image_1,
            "Image 2": image_2,
            "Image 3": image_3,
        }

        return io.NodeOutput(images[pick])


class ModelChiefExtension(ComfyExtension):

    @override
    async def get_node_list(self) -> list[type[io.ComfyNode]]:
        return [
            ImagePicker,
        ]


async def comfy_entrypoint() -> ModelChiefExtension:
    return ModelChiefExtension()