import nodes
from comfy_api.latest import ComfyExtension, io


class AnyNode(io.ComfyNode):
    ANYCLASS = True
    OUTPUTS = []
    INCLUDE = []
    EXCLUDE = []
    CATEGORY= "sampling/custom_sampler"

    @classmethod
    def define_schema(cls) -> io.Schema:
        combo_options = []
        for k in nodes.NODE_CLASS_MAPPINGS:
            klass = nodes.NODE_CLASS_MAPPINGS[k]
            if hasattr(klass, 'ANYCLASS'):
                continue
            if not hasattr(klass, 'define_schema'):
                continue
            if not all(x in k for x in cls.INCLUDE):
                continue
            if any(x in k for x in cls.EXCLUDE):
                continue
            s = klass.define_schema()
            if len(s.outputs) != len(cls.OUTPUTS) or not all(o1.__class__ == o2.__class__ for o1, o2 in zip(s.outputs, cls.OUTPUTS, strict=True)):
                continue
            combo_options.append(io.DynamicCombo.Option(k, s.inputs))
        return io.Schema(
                node_id=cls.__name__,
                display_name=cls.DISPLAY_NAME,
                category=cls.CATEGORY,
                outputs=cls.OUTPUTS,
                inputs=[io.DynamicCombo.Input("selection", options=combo_options)])


    @classmethod
    def execute(cls, selection):
        s = selection.pop("selection")
        klass = nodes.NODE_CLASS_MAPPINGS[s]
        
        return klass.execute(**selection)

class DynamicAnySampler(AnyNode):
    OUTPUTS = [io.Sampler.Output("sampler")] 
    INCLUDE = ["Sampler"]
    EXCLUDE = ["CustomSampler"]
    DISPLAY_NAME = "Any Sampler (All-In-One)"

class DynamicAnyGuider(AnyNode):
    OUTPUTS = [io.Guider.Output("guider")] 
    CATEGORY = "sampling/custom_sampling/"
    DISPLAY_NAME = "Any Guider (All-In-One)"

class DynamicAnyScheduler(AnyNode):
    OUTPUTS = [io.Sigmas.Output("sigmas")] 
    INCLUDE = ["Scheduler"]
    CATEGORY = "sampling/custom_sampling/"
    DISPLAY_NAME = "Any Scheduler (All-In-One)"

class DynamicAnyNoise(AnyNode):
    OUTPUTS = [io.Noise.Output("sigmas")] 
    INCLUDE = ["Noise"]
    CATEGORY = "sampling/custom_sampling/"
    DISPLAY_NAME = "Any Noise (All-In-One)"

class AnyNodeExtension(ComfyExtension):
    async def get_node_list(self):
        return [
            DynamicAnySampler,
            DynamicAnyGuider,
            DynamicAnyScheduler,
            DynamicAnyNoise,
        ]

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