import random

VARIATION_SLOTS = 10


class VariationBufferNode:

    OUTPUT_IS_LIST = (True,)

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {},
            "optional": {
                f"variation_{i}": ("STRING", {"forceInput": True})
                for i in range(1, VARIATION_SLOTS + 1)
            },
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("variations",)
    FUNCTION = "collect"
    CATEGORY = "conditioning/Character Prompt Manager"

    def collect(self, **kwargs):
        items = [
            kwargs[f"variation_{i}"]
            for i in range(1, VARIATION_SLOTS + 1)
            if kwargs.get(f"variation_{i}") is not None
        ]
        return (items,)


class CharacterPromptManagerNode:

    INPUT_IS_LIST = True
    OUTPUT_IS_LIST = (True, True)

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "base": ("STRING", {
                    "multiline": True,
                    "default": "masterpiece, best quality, vibe light",
                }),
                "character": ("STRING", {
                    "multiline": True,
                    "default": "old man sitting on a couch, grey hair, casual clothes",
                }),
                "others": ("STRING", {
                    "multiline": True,
                    "default": "simple background",
                }),
                "variations": ("STRING", {"forceInput": True}),
                "randomize_seed": ("BOOLEAN", {"default": True}),
            },
            "optional": {
                "manual_seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF}),
            },
        }

    RETURN_TYPES = ("STRING", "INT")
    RETURN_NAMES = ("text", "seed")
    FUNCTION = "generate"
    CATEGORY = "conditioning/Character Prompt Manager"

    @classmethod
    def IS_CHANGED(cls, randomize_seed, **kwargs):
        if randomize_seed[0]:
            return random.random()
        return False

    def generate(self, base, character, others, variations, randomize_seed, manual_seed=[0]):
        base = base[0]
        character = character[0]
        others = others[0]
        randomize_seed = randomize_seed[0]
        manual_seed = manual_seed[0] if manual_seed else 0

        lines = [v.strip() for v in variations if v.strip()]
        if not lines:
            return ([], [])

        texts = [f"{base},\n{character},\n{line},\n{others}" for line in lines]
        seeds = (
            [random.randint(0, 0xFFFFFFFF) for _ in lines]
            if randomize_seed
            else [manual_seed] * len(lines)
        )
        return (texts, seeds)


NODE_CLASS_MAPPINGS = {
    "CharacterPromptManager": CharacterPromptManagerNode,
    "VariationBuffer": VariationBufferNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "CharacterPromptManager": "Character Prompt Manager",
    "VariationBuffer": "Variation Buffer",
}