from __future__ import annotations

from comfy_api.latest import io


MSRCharacterData = io.Custom("MSR_CHARACTER")
MSRCharacterSetData = io.Custom("MSR_CHARACTER_SET")


def _compact_alias(alias: str) -> str:
    if not alias:
        return ""
    alias = alias.strip()
    if alias.startswith("@"):
        alias = alias[1:]
    return alias.strip()


class MSRCharacter(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="MSRCharacterCS",
            display_name="MSR Character CS",
            category="WhatDreamsCost CS",
            description=(
                "Reusable MSR character payload. Accepts up to two reference images plus "
                "a description and optional alias for @alias prompt replacement."
            ),
            inputs=[
                io.Image.Input(
                    "image_1",
                    optional=True,
                    tooltip="Optional first reference image for this character.",
                ),
                io.Image.Input(
                    "image_2",
                    optional=True,
                    tooltip="Optional second reference image for this character.",
                ),
                io.String.Input(
                    "description",
                    multiline=True,
                    default="",
                    optional=True,
                    tooltip="Manual or generated appearance description used for prompt replacement.",
                ),
                io.String.Input(
                    "alias",
                    default="",
                    optional=True,
                    tooltip="Optional alias. Prompts can use @alias alongside @charN/@characterN.",
                ),
            ],
            outputs=[
                MSRCharacterData.Output(display_name="character"),
            ],
        )

    @classmethod
    def execute(cls, image_1=None, image_2=None, description="", alias="") -> io.NodeOutput:
        images = []
        if image_1 is not None:
            images.append(image_1)
        if image_2 is not None:
            images.append(image_2)

        payload = {
            "images": images,
            "description": (description or "").strip(),
            "alias": _compact_alias(alias or ""),
        }
        return io.NodeOutput(payload)


class MSRCharacterSet(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="MSRCharacterSetCS",
            display_name="MSR Character Set CS",
            category="WhatDreamsCost CS",
            description="Bundles up to six MSR Character nodes into one character set for the director.",
            inputs=[
                MSRCharacterData.Input("character_1", optional=True),
                MSRCharacterData.Input("character_2", optional=True),
                MSRCharacterData.Input("character_3", optional=True),
                MSRCharacterData.Input("character_4", optional=True),
                MSRCharacterData.Input("character_5", optional=True),
                MSRCharacterData.Input("character_6", optional=True),
            ],
            outputs=[
                MSRCharacterSetData.Output(display_name="character_set"),
            ],
        )

    @classmethod
    def execute(
        cls,
        character_1=None,
        character_2=None,
        character_3=None,
        character_4=None,
        character_5=None,
        character_6=None,
    ) -> io.NodeOutput:
        characters = []
        for item in (
            character_1,
            character_2,
            character_3,
            character_4,
            character_5,
            character_6,
        ):
            if not item:
                continue
            characters.append(
                {
                    "images": list(item.get("images") or []),
                    "description": (item.get("description") or "").strip(),
                    "alias": _compact_alias(item.get("alias") or ""),
                }
            )

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