import uuid

from custom_nodes.comfyui_prompt_bench.api import OllamaChatMessage


class PBListChatView:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "chat": ("CHAT",),
                "text": (
                    "STRING",
                    {
                        "multiline": True,
                        "default": "",
                        "tooltip": "Latest rendered chat text.",
                    },
                ),
            },
            "optional": {}
        }

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        return str(uuid.uuid4())

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    FUNCTION = "process"
    CATEGORY = "Prompt Bench"
    OUTPUT_NODE = True

    def process(self, chat: list[OllamaChatMessage], text: str = ""):
        formatted_parts = []

        for msg in chat:
            role = msg.role.upper()
            content = msg.content.strip()
            formatted_parts.append(f"### {role} ###\n\n{content}\n\n---{'--------' if role == 'ASSISTANT' else ''}")

        text = "\n\n".join(formatted_parts)
        if not text:
            text = "(Empty Chat)"

        return {
            "ui": {"text": (text,)},
            "result": (text,),
        }