import uuid

from custom_nodes.comfyui_prompt_bench.api import _fetch_ollama_models, post_ollama_chat, OllamaChatMessage
from custom_nodes.comfyui_prompt_bench.prompt import build_prompt, refine_prompt, build_feedback

_SESSION_CHATS: dict[str, list[OllamaChatMessage]] = {}


class PBTextModelChat:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "session_seed": ("INT", {
                    "default": 0,
                    "min": 0,
                    "max": 0xffffffffffffffff,
                    "step": 1,
                    "tooltip": "Session key. Reuse to continue chat; change to start a new session.",
                }),
                "system_prompt": (
                    "STRING",
                    {
                        "multiline": True,
                        "default": "",
                        "dynamicPrompts": True,
                        "tooltip": "System prompt applied to the Ollama model.",
                    },
                ),
                "rating": ("RATING",),
                "model": (
                    "COMBO",
                    {
                        "options": _fetch_ollama_models(),
                        "remote": {
                            "route": "/prompt_bench/ollama/models",
                            "refresh_button": True,
                        },
                        "tooltip": "Ollama model name. Refresh to fetch the current local model list.",
                    },
                ),
            },
            "optional": {
                "text": ("STRING",),
            },
        }

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

    RETURN_TYPES = ("STRING", "STRING", "STRING", "CHAT",)
    RETURN_NAMES = ("enhanced_prompt", "model", "text", "chat",)
    FUNCTION = "process"
    CATEGORY = "Prompt Bench"

    def process(self, session_seed: int, system_prompt: str, rating: dict, model: str, text: str = ""):
        global _SESSION_CHATS

        session_id = str(session_seed)
        current_session = _SESSION_CHATS.get(session_id, [])
        has_session_messages = len(current_session) > 0
        previous_prompt = current_session[-1].content  if has_session_messages else ""

        prompt = build_prompt(
            previous_prompt if has_session_messages else "",
            "" if has_session_messages else system_prompt,
            build_feedback(rating, previous_prompt=previous_prompt),
            text
        )

        current_session.append(OllamaChatMessage(role="user", content=prompt))

        # Override here since we get feedback loop text
        assistant_message = post_ollama_chat(model, current_session)
        if assistant_message:
            current_session.append(assistant_message)

        _SESSION_CHATS[session_id] = current_session

        return assistant_message.content, model, text, _SESSION_CHATS[session_id],