"""
图像处理工具模块
提供 ComfyUI Tensor 与 PIL Image 之间的转换功能
"""

import base64
from io import BytesIO
from typing import List

import numpy as np
import torch
from PIL import Image


def tensor_to_pil(tensor: torch.Tensor) -> List[Image.Image]:
    """
    将 ComfyUI 的 Tensor 转换为 PIL Image 列表
    
    Args:
        tensor: 形状为 [B, H, W, C] 的张量，值范围 [0, 1]
    
    Returns:
        PIL Image 列表
    
    Example:
        >>> images = tensor_to_pil(input_tensor)
        >>> for img in images:
        ...     img.save(f"output_{i}.png")
    """
    images = []
    
    # 转换为 numpy 数组
    np_images = tensor.cpu().numpy()
    
    # 处理每张图像
    for i in range(np_images.shape[0]):
        img_array = np_images[i]
        
        # 转换值范围从 [0, 1] 到 [0, 255]
        img_array = (img_array * 255).astype(np.uint8)
        
        # 创建 PIL Image
        img = Image.fromarray(img_array)
        images.append(img)
    
    return images


def pil_to_tensor(images: List[Image.Image]) -> torch.Tensor:
    """
    将 PIL Image 列表转换为 ComfyUI 的 Tensor

    Args:
        images: PIL Image 列表

    Returns:
        形状为 [B, H, W, C] 的张量，值范围 [0, 1]

    Example:
        >>> pil_images = [Image.open("test.png")]
        >>> tensor = pil_to_tensor(pil_images)
        >>> print(tensor.shape)  # [1, H, W, 3]
    """
    tensors = []

    for img in images:
        # 确保是 RGB 模式
        if img.mode != 'RGB':
            img = img.convert('RGB')

        # 转换为 numpy 数组
        img_array = np.array(img).astype(np.float32)

        # 转换值范围从 [0, 255] 到 [0, 1]
        img_array = img_array / 255.0

        tensors.append(img_array)

    # 堆叠为批次
    batch_tensor = np.stack(tensors, axis=0)

    # 转换为 torch tensor
    return torch.from_numpy(batch_tensor)


def encode_image_to_base64(image: Image.Image, format: str = "PNG") -> str:
    """
    将 PIL Image 编码为 base64 字符串
    
    Args:
        image: PIL Image 对象
        format: 图像格式，默认 PNG
    
    Returns:
        base64 编码的字符串
    
    Example:
        >>> img = Image.open("test.png")
        >>> b64_str = encode_image_to_base64(img)
    """
    buffered = BytesIO()
    
    # 转换为 RGB 模式（如果是 RGBA）
    if image.mode == 'RGBA':
        image = image.convert('RGB')
    
    image.save(buffered, format=format)
    img_bytes = buffered.getvalue()
    
    return base64.b64encode(img_bytes).decode('utf-8')


_MAX_IMAGE_BYTES = 10 * 1024 * 1024  # 10MB base64 上限


def encode_image_to_base64_limited(
    image: Image.Image,
    format: str = "PNG",
    max_bytes: int = _MAX_IMAGE_BYTES,
) -> str:
    """
    将 PIL Image 编码为 base64，若超过 max_bytes 则自动缩放直到满足限制。

    策略：等比缩放，每轮缩小到上一轮的 80%，最多 10 轮。

    Args:
        image: PIL Image 对象
        format: 图像格式，默认 PNG
        max_bytes: base64 字符串最大字节数，默认 10MB

    Returns:
        base64 编码的字符串（保证 <= max_bytes）
    """
    working = image
    if working.mode == 'RGBA':
        working = working.convert('RGB')

    for attempt in range(10):
        buffered = BytesIO()
        working.save(buffered, format=format)
        b64 = base64.b64encode(buffered.getvalue()).decode('utf-8')

        if len(b64) <= max_bytes:
            if attempt > 0:
                print(
                    f"图片已自动缩放: {image.width}x{image.height} → "
                    f"{working.width}x{working.height} "
                    f"({len(b64) / 1024 / 1024:.2f}MB)"
                )
            return b64

        # 缩放到 80%
        scale = 0.8
        new_w = max(1, int(working.width * scale))
        new_h = max(1, int(working.height * scale))
        working = working.resize((new_w, new_h), Image.Resampling.LANCZOS)

    # 兜底：返回最后一次编码结果
    buffered = BytesIO()
    working.save(buffered, format=format)
    return base64.b64encode(buffered.getvalue()).decode('utf-8')


def decode_base64_to_pil(base64_string: str) -> Image.Image:
    """
    将 base64 字符串解码为 PIL Image
    
    Args:
        base64_string: base64 编码的图像字符串
    
    Returns:
        PIL Image 对象
    
    Example:
        >>> img = decode_base64_to_pil(b64_str)
        >>> img.save("decoded.png")
    """
    img_bytes = base64.b64decode(base64_string)
    img = Image.open(BytesIO(img_bytes))
    
    return img


def parse_batch_prompts(prompt: str) -> List[str]:
    """
    解析批量提示词
    
    检测单独行的 --- 分隔符，分割提示词。
    如果 --- 不是单独占据一行，则返回空列表（表示单提示词模式）。
    
    Args:
        prompt: 用户输入的提示词文本
    
    Returns:
        提示词列表。如果未检测到单独行的 ---，返回空列表（表示单提示词模式）
    
    Raises:
        ValueError: 如果所有提示词都为空
    
    Example:
        >>> prompts = parse_batch_prompts("a woman\\n---\\na man")
        >>> print(prompts)  # ['a woman', 'a man']
        
        >>> prompts = parse_batch_prompts("a woman --- a man")
        >>> print(prompts)  # []  (单提示词模式)
    """
    lines = prompt.split('\n')
    
    # 检查是否存在单独行的 ---
    has_separator = False
    for line in lines:
        if line.strip() == '---':
            has_separator = True
            break
    
    # 如果没有单独行的 ---，返回空列表（单提示词模式）
    if not has_separator:
        return []
    
    # 按单独行的 --- 分割
    # 先将所有单独行的 --- 替换为特殊标记
    processed_lines = []
    for line in lines:
        if line.strip() == '---':
            processed_lines.append('<<<SEPARATOR>>>')
        else:
            processed_lines.append(line)
    
    # 重新组合并分割
    processed_text = '\n'.join(processed_lines)
    raw_prompts = processed_text.split('<<<SEPARATOR>>>')
    
    # 过滤空提示词
    filtered_prompts = []
    for p in raw_prompts:
        stripped = p.strip()
        if stripped:
            filtered_prompts.append(stripped)
    
    # 如果所有提示词都为空，抛出错误
    if not filtered_prompts:
        raise ValueError("批量提示词模式下，所有提示词都为空，请至少提供一个有效的提示词")
    
    return filtered_prompts