from __future__ import annotations

from comfy_api.latest import io


TimelineSectionData = io.Custom("TIMELINE_SECTION")
TimelineSectionSetData = io.Custom("TIMELINE_SECTION_SET")
class TimelineSection(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="TimelineSectionCS",
            display_name="Timeline Section DUMAS",
            category="WhatDreamsCost DUMAS",
            description="Reusable timeline section with required text, optional image, and a duration in seconds.",
            inputs=[
                io.Image.Input(
                    "image",
                    optional=True,
                    tooltip="Optional image for this timeline section. When omitted, the section becomes a text-only block.",
                ),
                io.String.Input(
                    "text",
                    multiline=True,
                    default="",
                    tooltip="Required text/prompt for this timeline section.",
                ),
                io.Float.Input(
                    "duration",
                    default=1.0,
                    min=0.1,
                    max=1000.0,
                    step=0.01,
                    tooltip="Section duration in seconds.",
                ),
            ],
            outputs=[
                TimelineSectionData.Output(display_name="section"),
            ],
        )

    @classmethod
    def execute(cls, image=None, text="", duration=1.0) -> io.NodeOutput:
        payload = {
            "image": image,
            "text": (text or "").strip(),
            "duration": max(0.1, float(duration or 0.0)),
        }
        return io.NodeOutput(payload)


class TimelineSectionSet(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="TimelineSectionSetCS",
            display_name="Timeline Section Set DUMAS",
            category="WhatDreamsCost DUMAS",
            description="Bundles up to ten Timeline Section nodes into one reusable section set.",
            inputs=[
                TimelineSectionData.Input("section_1", optional=True),
                TimelineSectionData.Input("section_2", optional=True),
                TimelineSectionData.Input("section_3", optional=True),
                TimelineSectionData.Input("section_4", optional=True),
                TimelineSectionData.Input("section_5", optional=True),
                TimelineSectionData.Input("section_6", optional=True),
                TimelineSectionData.Input("section_7", optional=True),
                TimelineSectionData.Input("section_8", optional=True),
                TimelineSectionData.Input("section_9", optional=True),
                TimelineSectionData.Input("section_10", optional=True),
            ],
            outputs=[
                TimelineSectionSetData.Output(display_name="section_set"),
            ],
        )

    @classmethod
    def execute(
        cls,
        section_1=None,
        section_2=None,
        section_3=None,
        section_4=None,
        section_5=None,
        section_6=None,
        section_7=None,
        section_8=None,
        section_9=None,
        section_10=None,
    ) -> io.NodeOutput:
        sections = []
        for item in (
            section_1,
            section_2,
            section_3,
            section_4,
            section_5,
            section_6,
            section_7,
            section_8,
            section_9,
            section_10,
        ):
            if not item:
                continue
            text = (item.get("text") or "").strip()
            if not text:
                continue
            sections.append(
                {
                    "image": item.get("image"),
                    "text": text,
                    "duration": max(0.1, float(item.get("duration") or 0.0)),
                }
            )

        return io.NodeOutput({"sections": sections})