import base64
import io
from PIL import Image
import numpy as np
from openai import OpenAI


class ImageToLLMPrompt:
    """
    A custom node that uploads an image and sends it to an OpenAI-compatible API
    with optional system and user prompts, returning the LLM response as text.
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "base_url": ("STRING", {
                    "default": "http://localhost:8080/v1",
                    "multiline": False
                }),
                "api_key": ("STRING", {
                    "default": "sk-test",
                    "multiline": False
                }),
                "model": ("STRING", {
                    "default": "gpt-4o",
                    "multiline": False
                }),
            },
            "optional": {
                "system_prompt": ("STRING", {
                    "default": "",
                    "multiline": True
                }),
                "user_prompt": ("STRING", {
                    "default": "",
                    "multiline": True
                }),
                "temperature": ("FLOAT", {
                    "default": 0.7,
                    "min": 0.0,
                    "max": 2.0,
                    "step": 0.1
                }),
                "max_tokens": ("INT", {
                    "default": 2048,
                    "min": 1,
                    "max": 32768,
                    "step": 1
                }),
            }
        }
    
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    FUNCTION = "process_image"
    CATEGORY = "LLM"
    OUTPUT_NODE = True
    
    def tensor_to_base64(self, image_tensor):
        """Convert ComfyUI image tensor to base64 string"""
        # ComfyUI images are in format [B, H, W, C] with values 0-1
        image_np = (image_tensor[0].cpu().numpy() * 255).astype(np.uint8)
        
        # Convert to PIL Image
        pil_image = Image.fromarray(image_np)
        
        # Convert to base64
        buffer = io.BytesIO()
        pil_image.save(buffer, format="PNG")
        buffer.seek(0)
        base64_image = base64.b64encode(buffer.read()).decode('utf-8')
        
        return base64_image
    
    def process_image(self, image, base_url, api_key, model, 
                     system_prompt="", user_prompt="", temperature=0.7, max_tokens=2048):
        """Send image to LLM and return response"""
        
        # Convert image to base64
        base64_image = self.tensor_to_base64(image)
        
        # Create OpenAI client
        client = OpenAI(
            base_url=base_url,
            api_key=api_key,
            timeout=180
        )
        
        # Build messages array
        messages = []
        
        # Add optional system prompt
        if system_prompt and system_prompt.strip():
            messages.append({
                "role": "system",
                "content": system_prompt
            })
        
        # Build user message content
        user_content = []
        
        # Add optional user prompt text
        if user_prompt and user_prompt.strip():
            user_content.append({
                "type": "text",
                "text": user_prompt
            })
        
        # Add the image
        user_content.append({
            "type": "image_url",
            "image_url": {
                "url": f"data:image/png;base64,{base64_image}"
            }
        })
        
        messages.append({
            "role": "user",
            "content": user_content
        })
        
        # Make API call
        try:
            print(f"[ImageToLLMPrompt] Sending request to {base_url} with model {model}")
            response = client.chat.completions.create(
                model=model,
                temperature=temperature,
                max_tokens=max_tokens,
                messages=messages
            )
            
            result = response.choices[0].message.content
            print(f"[ImageToLLMPrompt] LLM Response received (length: {len(result)} chars)")
            print(f"[ImageToLLMPrompt] Response content:\n{result}")
            return {"ui": {"text": (result,)}, "result": (result,)}
            
        except Exception as e:
            error_msg = f"LLM API call failed: {str(e)}"
            print(f"[ImageToLLMPrompt] ERROR: {error_msg}")
            return {"ui": {"text": (error_msg,)}, "result": (error_msg,)}


# Node registration
NODE_CLASS_MAPPINGS = {
    "ImageToLLMPrompt": ImageToLLMPrompt
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ImageToLLMPrompt": "Image to LLM Prompt"
}