import base64
import json
import requests
import numpy as np
from PIL import Image

class LMStudioVisionPrompt:
    """
    Custom ComfyUI node:
    - Input: IMAGE
    - Calls LM Studio vision model (OpenAI-style /v1/chat/completions)
    - Output: single-line SD-style prompt (STRING)
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
            },
            "optional": {
                # NOTE: uses your LM Studio address
                "lmstudio_url": (
                    "STRING",
                    {
                        "default": "http://127.0.0.1:1234/v1/chat/completions",
                        "multiline": False,
                    },
                ),
                "model_name": (
                    "STRING",
                    {
                        "default": "your-vision-model-name",
                        "multiline": False,
                    },
                ),
                "extra_instruction": (
                    "STRING",
                    {
                        "default": (
                            "Describe this image as a single-line Stable Diffusion prompt. "
                            "Include subject, pose, environment, lighting and mood. "
                            "Do NOT explain, just output the prompt."
                        ),
                        "multiline": True,
                    },
                ),
                "max_tokens": ("INT", {"default": 150, "min": 16, "max": 512}),
            },
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("prompt",)
    FUNCTION = "run"
    CATEGORY = "LM Studio"

    def _image_to_base64(self, image_tensor):
        """
        Convert ComfyUI IMAGE tensor to base64 PNG string.
        image_tensor: (B, H, W, C) in [0,1]
        """
        # Use the first image in batch
        img = image_tensor[0].cpu().numpy()  # (H, W, C)
        img = np.clip(img * 255.0, 0, 255).astype(np.uint8)
        pil_img = Image.fromarray(img)

        from io import BytesIO
        buffer = BytesIO()
        pil_img.save(buffer, format="PNG")
        b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
        return b64

    def run(
        self,
        image,
        lmstudio_url,
        model_name,
        extra_instruction,
        max_tokens,
    ):
        # Convert image to base64
        image_b64 = self._image_to_base64(image)

        # Build request payload for LM Studio (OpenAI-compatible)
        payload = {
            "model": model_name,
            "messages": [
                {
                    "role": "system",
                    "content": "You are a Stable Diffusion prompt generator.",
                },
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": extra_instruction,
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{image_b64}"
                            },
                        },
                    ],
                },
            ],
            "max_tokens": max_tokens,
        }

        try:
            response = requests.post(
                lmstudio_url,
                headers={"Content-Type": "application/json"},
                data=json.dumps(payload),
                timeout=60,
            )
            response.raise_for_status()
            data = response.json()

            prompt_text = data["choices"][0]["message"]["content"].strip()

        except Exception as e:
            # If something goes wrong, return an error string as the prompt
            prompt_text = f"[LM Studio error: {e}]"

        return (prompt_text,)


# Required ComfyUI registration
NODE_CLASS_MAPPINGS = {
    "LMStudioVisionPrompt": LMStudioVisionPrompt,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "LMStudioVisionPrompt": "LM Studio Vision → Prompt",
}