"""
BoosAPI ComfyUI Custom Nodes

让 ComfyUI workflow 调用 BoosAPI 的 AI 模型（图像生成、LLM 等）。
配置：通过环境变量设置 BoosAPI 地址，或使用默认 http://127.0.0.1:4000/v1

环境变量：
  BOOSAPI_BASE    — BoosAPI 地址（默认 http://127.0.0.1:4000）
  BOOSAPI_KEY     — API Key（默认 sk-boos4721）
"""

import json
import os
import requests

# ── Config ───────────────────────────────────────────────────────────────
BOOSAPI_BASE = os.environ.get("BOOSAPI_BASE", "http://127.0.0.1:4000")
BOOSAPI_KEY = os.environ.get("BOOSAPI_KEY", "sk-boos4721")
HEADERS = {
    "Authorization": f"Bearer {BOOSAPI_KEY}",
    "Content-Type": "application/json",
}

# ── Node Mappings ────────────────────────────────────────────────────────
NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}

# ═════════════════════════════════════════════════════════════════════════
# Node: BoosAPIImageGen — 图像生成 (OpenAI DALL-E / 兼容格式)
# ═════════════════════════════════════════════════════════════════════════

class BoosAPIImageGen:
    """调用 BoosAPI /v1/images/generations 生成图像"""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "prompt": ("STRING", {"multiline": True, "default": ""}),
                "model": (["dall-e-3", "dall-e-2", "gpt-image-2"], {"default": "gpt-image-2"}),
                "size": (["1024x1024", "1792x1024", "1024x1792"], {"default": "1024x1024"}),
                "quality": (["standard", "hd"], {"default": "standard"}),
                "n": ("INT", {"default": 1, "min": 1, "max": 4}),
            },
            "optional": {
                "style": (["vivid", "natural"], {"default": "vivid"}),
            },
        }

    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("image", "metadata")
    FUNCTION = "generate"
    CATEGORY = "BoosAPI"

    def generate(self, prompt, model, size, quality, n, style="vivid"):
        url = f"{BOOSAPI_BASE}/v1/images/generations"
        payload = {
            "model": model,
            "prompt": prompt,
            "n": n,
            "size": size,
            "quality": quality,
            "style": style,
            "response_format": "b64_json",
        }
        try:
            resp = requests.post(url, headers=HEADERS, json=payload, timeout=120)
            resp.raise_for_status()
            data = resp.json()
            images = data.get("data", [])
            # Return first image as b64
            from io import BytesIO
            import base64
            import torch
            from PIL import Image

            b64_data = images[0]["b64_json"]
            img_bytes = base64.b64decode(b64_data)
            pil_img = Image.open(BytesIO(img_bytes)).convert("RGB")
            img_tensor = torch.from_numpy(
                __import__("numpy").array(pil_img).astype("float32") / 255.0
            ).unsqueeze(0)

            metadata = json.dumps({
                "model": model,
                "size": size,
                "revised_prompt": data.get("data", [{}])[0].get("revised_prompt", ""),
            })
            return (img_tensor, metadata)
        except Exception as e:
            print(f"[BoosAPI] ImageGen failed: {e}")
            raise


# ═════════════════════════════════════════════════════════════════════════
# Node: BoosAPILLM — LLM Chat (流式文本生成)
# ═════════════════════════════════════════════════════════════════════════

class BoosAPILLM:
    """调用 BoosAPI /v1/chat/completions 生成文本"""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "prompt": ("STRING", {"multiline": True, "default": ""}),
                "model": ("STRING", {"default": "gpt-5.5"}),
                "system_prompt": ("STRING", {"multiline": True, "default": ""}),
                "temperature": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 2.0}),
                "max_tokens": ("INT", {"default": 2048, "min": 1, "max": 16384}),
            },
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    FUNCTION = "chat"
    CATEGORY = "BoosAPI"

    def chat(self, prompt, model, system_prompt, temperature, max_tokens):
        url = f"{BOOSAPI_BASE}/v1/chat/completions"
        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": prompt})

        payload = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
            "stream": False,
        }
        try:
            resp = requests.post(url, headers=HEADERS, json=payload, timeout=120)
            resp.raise_for_status()
            data = resp.json()
            text = data["choices"][0]["message"]["content"]
            return (text,)
        except Exception as e:
            print(f"[BoosAPI] LLM chat failed: {e}")
            raise


# ═════════════════════════════════════════════════════════════════════════
# Node: BoosAPITextToImage — 文生图 (ComfyUI 原生格式)
# ═════════════════════════════════════════════════════════════════════════

class BoosAPITextToImage:
    """输入文本提示词，通过 BoosAPI 生成图像，输出 ComfyUI IMAGE 格式"""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "positive": ("STRING", {"multiline": True, "default": ""}),
                "model": (["gpt-image-2", "dall-e-3"], {"default": "gpt-image-2"}),
            },
            "optional": {
                "negative": ("STRING", {"multiline": True, "default": ""}),
                "size": (["1024x1024", "1792x1024", "1024x1792"], {"default": "1024x1024"}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "generate"
    CATEGORY = "BoosAPI"

    def generate(self, positive, model, negative="", size="1024x1024"):
        # Build prompt with negative guidance
        prompt = positive
        if negative:
            prompt = f"{positive}\n避免: {negative}"

        url = f"{BOOSAPI_BASE}/v1/images/generations"
        payload = {
            "model": model,
            "prompt": prompt,
            "n": 1,
            "size": size,
            "response_format": "b64_json",
        }
        try:
            resp = requests.post(url, headers=HEADERS, json=payload, timeout=120)
            resp.raise_for_status()
            data = resp.json()

            import base64
            from io import BytesIO
            import torch
            from PIL import Image
            import numpy as np

            b64 = data["data"][0]["b64_json"]
            img = Image.open(BytesIO(base64.b64decode(b64))).convert("RGB")
            tensor = torch.from_numpy(np.array(img).astype("float32") / 255.0).unsqueeze(0)
            return (tensor,)
        except Exception as e:
            print(f"[BoosAPI] T2I failed: {e}")
            raise


# ═════════════════════════════════════════════════════════════════════════
# Register Nodes
# ═════════════════════════════════════════════════════════════════════════

NODE_CLASS_MAPPINGS = {
    "BoosAPIImageGen": BoosAPIImageGen,
    "BoosAPILLM": BoosAPILLM,
    "BoosAPITextToImage": BoosAPITextToImage,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "BoosAPIImageGen": "BoosAPI 图像生成",
    "BoosAPILLM": "BoosAPI LLM 对话",
    "BoosAPITextToImage": "BoosAPI 文生图",
}

__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]