from __future__ import annotations

from comfy_api.latest import io


TimelineSectionData = io.Custom("TIMELINE_SECTION")
TimelineSectionSetData = io.Custom("TIMELINE_SECTION_SET")
TimelineData = io.Custom("TIMELINE")


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 five 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),
            ],
            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,
    ) -> io.NodeOutput:
        sections = []
        for item in (section_1, section_2, section_3, section_4, section_5):
            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})


class Timeline(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="TimelineCS",
            display_name="Timeline DUMAS",
            category="WhatDreamsCost DUMAS",
            description="Builds an external timeline payload for LTX Director from a section set, deriving the total duration from the sections.",
            inputs=[
                TimelineSectionSetData.Input(
                    "section_set",
                    optional=True,
                    tooltip="Optional section set. When omitted, the timeline payload contains no sections.",
                ),
            ],
            outputs=[
                TimelineData.Output(display_name="timeline"),
            ],
        )

    @classmethod
    def execute(cls, section_set=None, duration=None) -> io.NodeOutput:
        sections = list((section_set or {}).get("sections") or [])
        total_duration = sum(max(0.1, float(item.get("duration") or 0.0)) for item in sections)
        payload = {
            "duration": max(0.1, float(total_duration or 0.0)),
            "sections": sections,
        }
        return io.NodeOutput(payload)