"""Underpin Optimise text node. Save as custom_nodes/underpin_optimise/__init__.py."""
import json
import os
import urllib.request
import urllib.error


class UnderpinOptimise:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {
            "text": ("STRING", {"multiline": True}),
            "query": ("STRING", {"multiline": True}),
            "ratio": ("INT", {"default": 5, "min": 1, "max": 20}),
        }}

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("context",)
    FUNCTION = "optimise"
    CATEGORY = "Underpin"

    def optimise(self, text, query, ratio):
        key = os.environ.get("OPTIMISE_API_KEY")
        if not key:
            raise RuntimeError("Set OPTIMISE_API_KEY in the environment that starts ComfyUI, then restart it.")
        request = urllib.request.Request(
            "https://underpin-optimise-api.westernroseglobal.workers.dev/v1/compress",
            data=json.dumps({"text": text, "query": query, "ratio": ratio, "guarded": True}).encode(),
            headers={"Content-Type": "application/json", "Authorization": "Bearer " + key,
                     "User-Agent": "underpin-comfyui/1.0"},
        )
        try:
            with urllib.request.urlopen(request, timeout=20) as response:
                result = json.load(response)
        except urllib.error.HTTPError as error:
            raise RuntimeError("Optimise HTTP " + str(error.code) + ": " + error.read().decode()) from error
        except urllib.error.URLError as error:
            raise RuntimeError("Optimise could not connect. Check the network and account usage before retrying.") from error
        if not isinstance(result.get("text"), str):
            raise RuntimeError("Optimise returned no text.")
        return (result["text"],)


NODE_CLASS_MAPPINGS = {"UnderpinOptimise": UnderpinOptimise}
NODE_DISPLAY_NAME_MAPPINGS = {"UnderpinOptimise": "Underpin Optimise — Text Context"}