import json
import uuid

from custom_nodes.comfyui_prompt_bench.api import _fetch_ollama_models, post_ollama_text
from custom_nodes.comfyui_prompt_bench.prompt import build_prompt


class PBTextModel:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "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",)
    RETURN_NAMES = ("enhanced_prompt", "model", "text",)
    FUNCTION = "process"
    CATEGORY = "Prompt Bench"

    def process(self, system_prompt: str, rating: dict, model: str, text: str = ""):
        rating_context = json.dumps(rating, ensure_ascii=True, sort_keys=True)
        prompt = build_prompt("", system_prompt, rating_context, text)

        # Override here since we get feedback loop text
        text = post_ollama_text(model, prompt)

        return text, model, text,