import importlib
parse_prompt_schedules = importlib.import_module("custom_nodes.comfyui-prompt-control.prompt_control.parser").parse_prompt_schedules

class PCExtractScheduledPrompt
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "text": ("STRING", {"multiline": True}),
                "timestep": ("FLOAT", {"min": 0.0, "max": 1.0, "default": 1.0, "step": 0.01}),
            }
        }

    RETURN_TYPES = ("STRING",)
    CATEGORY = "promptcontrol/tools"
    FUNCTION = "apply"

    def apply(self, text, timestep):
        schedule = parse_prompt_schedules(text)
        _, entry = schedule.at_step(timestep, total_steps=1)
        prompt_text = entry.get("prompt", "")
        return (prompt_text,)

NODE_CLASS_MAPPINGS = {
    "PCExtractScheduledPrompt": PCExtractScheduledPrompt
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "PCExtractScheduledPrompt": "PC: Extract Scheduled Prompt",
}