# ComfyUI 节点：api_caption（精简版）
import base64
import io
import json
from typing import Tuple

import numpy as np
import requests
from PIL import Image
import torch


def tensor2pil(t_image: torch.Tensor) -> Image.Image:
    if not isinstance(t_image, torch.Tensor):
        raise TypeError("输入必须为 torch.Tensor")
    arr = t_image.cpu().numpy()
    if arr.ndim == 4:
        arr = arr[0]
    arr = np.clip(255.0 * arr.squeeze(), 0, 255).astype(np.uint8)
    if arr.ndim == 3 and arr.shape[0] in (1, 3):
        arr = np.transpose(arr, (1, 2, 0))
    return Image.fromarray(arr)


class api_caption:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "API类型": (["硅基流动", "贞贞工坊", "OpenRouter", "其他"], {"default": "硅基流动"}),
                "请求地址": ("STRING", {"default": "<url>"}),
                "API_Key": ("STRING", {"default": "<your_key>"}),
                "模型名称": ("STRING", {"default": "Qwen/Qwen3-VL-32B-Instruct"}),
                "image": ("IMAGE",),
                "提示词": ("STRING", {"default": "请用自然语言描述该图片。直接返回描述，不要有其他废话。字数控制在50字以内。", "multiline": True, "rows": 4}),
                "输出语言": (["中文", "英文"], {"default": "中文"}),
                "温度": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 2.0, "step": 0.01}),
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    FUNCTION = "generate"
    CATEGORY = "DaNodes/API"

    def generate(self, API类型: str, API_Key: str, 模型名称: str, image: torch.Tensor, 提示词: str,
                 输出语言: str, 温度: float = 0.5, 请求地址: str = "") -> Tuple[str]:
        try:
            # 根据 provider 选择实际的 URL
            provider_map = {
                "OpenRouter": "https://openrouter.ai/api/v1/chat/completions",
                "硅基流动": "https://api.siliconflow.cn/v1/chat/completions",
                "贞贞工坊": "https://api.bltcy.ai/v1/chat/completions",
            }
            if API类型 == "其他":
                请求地址 = 请求地址 or ""
            else:
                请求地址 = provider_map.get(API类型, 请求地址)
            if not 请求地址:
                return ("请求地址为空，请填写或选择有效的 API 类型",)
            # image -> base64（保持原始尺寸，使用无损 PNG 编码）
            pil = tensor2pil(image)
            buf = io.BytesIO()
            pil.save(buf, format="PNG")
            img_b64 = base64.b64encode(buf.getvalue()).decode()

            # 提示词
            if 输出语言 == "中文":
                提示词_full = f"{提示词} 请用中文返回描述"
            else:
                提示词_full = f"{提示词} Please return the description in English."

            # 按照示例将图片放到 messages[0].content 中
            image_data_uri = f"data:image/png;base64,{img_b64}"
            messages = [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": 提示词_full},
                        {"type": "image_url", "image_url": {"url": image_data_uri}},
                    ],
                }
            ]

            payload = {"model": 模型名称, "stream": False, "messages": messages, "max_tokens": 400, "temperature": float(温度)}
            headers = {
                "Authorization": f"Bearer {API_Key}",
                "Content-Type": "application/json",
                "Accept": "application/json",
                "User-Agent": "ComfyUI-DaNodes/1.0",
            }

            # 固定使用硅基流动的 JSON 风格请求
            try:
                r = requests.post(请求地址, json=payload, headers=headers, timeout=30)
            except Exception as e:
                return (f"API 请求失败（网络/连接错误）: {e}",)

            if not r.ok:
                body = None
                try:
                    body = r.json()
                    body = json.dumps(body, ensure_ascii=False)
                except Exception:
                    body = r.text
                return (f"API 返回错误状态 {r.status_code}: {body}",)

            try:
                result = r.json()
            except Exception as e:
                return (f"无法解析 API 返回为 JSON: {e}; body: {r.text}",)

            # 提取文本回复的帮助函数
            def extract_text_from_content(content):
                if isinstance(content, str):
                    return content
                if isinstance(content, dict):
                    return content.get("text") or content.get("content") or str(content)
                if isinstance(content, list):
                    parts = []
                    for block in content:
                        if not isinstance(block, dict):
                            parts.append(str(block))
                            continue
                        btype = block.get("type", "")
                        if btype in ("text", "paragraph"):
                            parts.append(block.get("text", ""))
                        elif btype == "image_url":
                            continue
                        else:
                            parts.append(block.get("text") or block.get("content") or "")
                    return " ".join(p for p in parts if p).strip()
                return str(content)

            caption = ""
            if isinstance(result, dict):
                choices = result.get("choices") or []
                if choices:
                    first = choices[0]
                    if isinstance(first, dict):
                        msg = first.get("message") or first.get("delta") or None
                        if isinstance(msg, dict) and "content" in msg:
                            caption = extract_text_from_content(msg.get("content"))
                        else:
                            if "text" in first and isinstance(first.get("text"), str):
                                caption = first.get("text")
                            else:
                                caption = extract_text_from_content(first)
                    else:
                        caption = str(first)
                else:
                    caption = result.get("caption") or result.get("text") or ""
            else:
                caption = str(result)

            caption = caption.strip()
            if caption == "":
                try:
                    pretty = json.dumps(result, ensure_ascii=False)
                except Exception:
                    pretty = str(result)
                return (f"无法从 API 响应解析出文本。返回 JSON: {pretty}",)

            return (caption,)
        except Exception as e:
            # 捕获整个函数体的未处理异常，避免 ComfyUI 无输出
            return (f"节点内部异常: {e}",)

NODE_CLASS_MAPPINGS = {
    "api_caption": api_caption
}
NODE_DISPLAY_NAME_MAPPINGS = {
    "api_caption": "API 打标（精简版）☀"
}