import base64
import io
from typing import Dict, Any, List, Tuple, Union

import torch
from PIL import Image
import requests
from requests.adapters import HTTPAdapter
from requests.packages.urllib3.util.retry import Retry


class Image2Prompt:
    """
    ComfyUI 自定义节点：将图片转换为文本提示
    """

    @classmethod
    def INPUT_TYPES(cls) -> Dict[str, Any]:
        return {
            "required": {
                "images": ("IMAGE",),
                "APIUrl": ("STRING", {"default": "", "multiline": False}),
                "APIKey": ("STRING", {"default": "", "multiline": False}),
                "RetryCount": ("INT", {"default": 3, "min": 0, "max": 100}),
                "Neutral": ("BOOLEAN", {"default": True}),
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    FUNCTION = "execute"
    CATEGORY = "NudeJS"

    def __init__(self):
        self.session = requests.Session()
        # 配置重试策略
        retry_strategy = Retry(
            total=3,
            backoff_factor=1,
            status_forcelist=[429, 500, 502, 503, 504],
        )
        adapter = HTTPAdapter(max_retries=retry_strategy)
        self.session.mount("http://", adapter)
        self.session.mount("https://", adapter)

    def execute(
        self,
        images: torch.Tensor,
        APIUrl: str,
        APIKey: str,
        RetryCount: int,
        Neutral: bool,
    ) -> Tuple[Union[str, None]]:
        """
        执行图片转提示词的主逻辑

        Args:
            images: 输入的图片张量
            APIUrl: API 接口地址
            APIKey: API 密钥
            RetryCount: 重试次数
            Neutral: 是否使用中性文本

        Returns:
            转换后的文本提示
        """
        if len(images) == 0:
            return (None,)

        # 获取第一张图片（ComfyUI 通常批量处理，这里取第一张）
        img_tensor = images[0]

        # 转换为 PIL Image
        pil_image = self._tensor_to_pil(img_tensor)

        # 检查并缩放图片
        resized_image = self._resize_image(pil_image)

        # 转换为 base64
        img_base64 = self._image_to_base64(resized_image)

        # 发送请求获取提示词
        prompt_text = self._call_api(
            APIUrl=APIUrl,
            APIKey=APIKey,
            imgBase64=img_base64,
            retry_count=RetryCount,
            neutral=Neutral,
        )

        return (prompt_text,)

    def _tensor_to_pil(self, img_tensor: torch.Tensor) -> Image.Image:
        """
        将 ComfyUI 的张量转换为 PIL Image

        Args:
            img_tensor: ComfyUI 中的图片张量 (B, H, W, C) 或 (H, W, C)

        Returns:
            PIL Image
        """
        # 确保张量在 CPU 上
        if img_tensor.is_cuda:
            img_tensor = img_tensor.cpu()

        # 处理不同的张量形状
        if len(img_tensor.shape) == 4:
            # 批量图片，取第一张
            img_tensor = img_tensor[0]
        elif len(img_tensor.shape) == 3:
            # 单张图片
            pass
        else:
            raise ValueError(f"Unexpected tensor shape: {img_tensor.shape}")

        # 转换为 PIL Image
        # ComfyUI 的图片张量通常值范围是 [0, 1]，形状是 [H, W, C]
        if img_tensor.shape[-1] in [3, 4]:  # RGB 或 RGBA
            array = (img_tensor.numpy() * 255).astype("uint8")
            pil_image = Image.fromarray(array, mode="RGB")
        else:
            raise ValueError(f"Unexpected number of channels: {img_tensor.shape[-1]}")

        return pil_image

    def _resize_image(self, image: Image.Image) -> Image.Image:
        """
        缩放图片，确保宽度不超过 1000，高度不超过 800

        Args:
            image: 输入的 PIL 图片

        Returns:
            缩放后的图片
        """
        width, height = image.size
        new_width = width
        new_height = height

        # 检查宽度
        if width > 1000:
            new_width = 1000
            new_height = int((height * 1000) / width)

        # 检查高度
        if new_height > 800:
            new_height = 800
            new_width = int((new_width * 800) / new_height)

        # 如果需要缩放，则调整大小
        if new_width != width or new_height != height:
            image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)

        return image

    def _image_to_base64(self, image: Image.Image) -> str:
        """
        将 PIL 图片转换为 base64 字符串（不包含 data:image/png;base64, 前缀）

        Args:
            image: 输入的 PIL 图片

        Returns:
            base64 编码的字符串
        """
        # 转换为 RGB 模式（如果需要）
        if image.mode != "RGB":
            image = image.convert("RGB")

        # 保存到内存中的字节流
        buffer = io.BytesIO()
        image.save(buffer, format="PNG")
        img_bytes = buffer.getvalue()

        # 编码为 base64
        img_base64 = base64.b64encode(img_bytes).decode("utf-8")

        return img_base64

    def _call_api(
        self,
        APIUrl: str,
        APIKey: str,
        imgBase64: str,
        retry_count: int,
        neutral: bool,
    ) -> Union[str, None]:
        """
        调用 API 接口获取提示词

        Args:
            APIUrl: API 地址
            APIKey: API 密钥
            imgBase64: base64 编码的图片
            retry_count: 最大重试次数
            neutral: 是否使用中性文本

        Returns:
            提示词文本，失败返回 None
        """
        url = f"{APIUrl}/generate"
        headers = {"authorization": APIKey}
        data = { "type": "img2img", "base64": imgBase64 }

        max_retries = max(0, retry_count)

        for attempt in range(max_retries + 1):
            try:
                response = self.session.post(url, json=data, headers=headers, timeout=30)
                response.raise_for_status()

                result = response.json()

                if result.get("code") == 200:
                    if neutral:
                        return result.get("neutralText")
                    else:
                        return result.get("promptText")
                else:
                    # API 返回错误码，记录日志
                    print(f"API request failed with code: {result.get('code')}, message: {result.get('message', 'unknown error')}")

            except requests.exceptions.RequestException as e:
                print(f"Attempt {attempt + 1} failed: {str(e)}")

            except Exception as e:
                print(f"Unexpected error on attempt {attempt + 1}: {str(e)}")

            # 如果不是最后一次尝试，等待一下再重试
            if attempt < max_retries:
                import time
                time.sleep(1)

        # 所有重试都失败了
        print(f"Failed after {max_retries} retries")
        return None


class Text2Prompt:
    """
    ComfyUI 自定义节点：将文本描述转换为图片提示
    """

    @classmethod
    def INPUT_TYPES(cls) -> Dict[str, Any]:
        return {
            "required": {
                "imgDesc": ("STRING", {"default": "", "multiline": True}),
                "APIUrl": ("STRING", {"default": "", "multiline": False}),
                "APIKey": ("STRING", {"default": "", "multiline": False}),
                "RetryCount": ("INT", {"default": 3, "min": 0, "max": 100}),
                "Neutral": ("BOOLEAN", {"default": True}),
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    FUNCTION = "execute"
    CATEGORY = "NudeJS"

    def __init__(self):
        self.session = requests.Session()
        # 配置重试策略
        retry_strategy = Retry(
            total=3,
            backoff_factor=1,
            status_forcelist=[429, 500, 502, 503, 504],
        )
        adapter = HTTPAdapter(max_retries=retry_strategy)
        self.session.mount("http://", adapter)
        self.session.mount("https://", adapter)

    def execute(
        self,
        imgDesc: str,
        APIUrl: str,
        APIKey: str,
        RetryCount: int,
        Neutral: bool,
    ) -> Tuple[Union[str, None]]:
        """
        执行文本转提示词的主逻辑

        Args:
            imgDesc: 图片的文本描述
            APIUrl: API 接口地址
            APIKey: API 密钥
            RetryCount: 重试次数
            Neutral: 是否使用中性文本

        Returns:
            转换后的文本提示
        """

        # 发送请求获取提示词
        prompt_text = self._call_api(
            APIUrl=APIUrl,
            APIKey=APIKey,
            imgDesc=imgDesc,
            retry_count=RetryCount,
            neutral=Neutral,
        )

        return (prompt_text,)

    def _call_api(
        self,
        APIUrl: str,
        APIKey: str,
        imgDesc: str,
        retry_count: int,
        neutral: bool,
    ) -> Union[str, None]:
        """
        调用 API 接口获取提示词

        Args:
            APIUrl: API 地址
            APIKey: API 密钥
            imgDesc: 图片的文本描述
            retry_count: 最大重试次数
            neutral: 是否使用中性文本

        Returns:
            提示词文本，失败返回 None
        """
        url = f"{APIUrl}/generate"
        headers = {"authorization": APIKey}
        data = { "type": "text2img", "desc": imgDesc }

        max_retries = max(0, retry_count)

        for attempt in range(max_retries + 1):
            try:
                response = self.session.post(url, json=data, headers=headers, timeout=30)
                response.raise_for_status()

                result = response.json()

                if result.get("code") == 200:
                    if neutral:
                        return result.get("neutralText")
                    else:
                        return result.get("promptText")
                else:
                    # API 返回错误码，记录日志
                    print(f"API request failed with code: {result.get('code')}, message: {result.get('message', 'unknown error')}")

            except requests.exceptions.RequestException as e:
                print(f"Attempt {attempt + 1} failed: {str(e)}")

            except Exception as e:
                print(f"Unexpected error on attempt {attempt + 1}: {str(e)}")

            # 如果不是最后一次尝试，等待一下再重试
            if attempt < max_retries:
                import time
                time.sleep(1)

        # 所有重试都失败了
        print(f"Failed after {max_retries} retries")
        return None


# 注册节点
NODE_CLASS_MAPPINGS = {
    "Image2Prompt": Image2Prompt,
    "Text2Prompt": Text2Prompt,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "Image2Prompt": "Image2Prompt",
    "Text2Prompt": "Text2Prompt",
}