ForcePilot/backend/package/yuxi/utils/image_processor.py

"""
图片处理工具模块
支持图片的格式转换、压缩、缩略图生成等功能
"""

import base64
import io

from PIL import ExifTags, Image

from yuxi.utils import logger


class ImageProcessor:
    """图片处理类"""

    # 支持的图片格式
    SUPPORTED_FORMATS = {"JPEG", "PNG", "WebP", "GIF", "BMP"}

    # 最大文件大小（5MB）
    MAX_FILE_SIZE = 5 * 1024 * 1024

    # 缩略图尺寸
    THUMBNAIL_SIZE = (200, 200)

    def process_image(self, image_data: bytes, original_filename: str = "") -> dict:
        """
        处理上传的图片

        Args:
            image_data: 图片二进制数据
            original_filename: 原始文件名

        Returns:
            dict: 包含处理结果的字典
        """
        try:
            # 验证图片格式
            img_format, _ = self._validate_image_format(image_data)
            if img_format not in self.SUPPORTED_FORMATS:
                raise ValueError(f"不支持的图片格式: {img_format}")

            # 加载图片
            with Image.open(io.BytesIO(image_data)) as img:
                # 处理EXIF方向信息
                img = self._fix_image_orientation(img)

                # 生成缩略图
                thumbnail_data = self._generate_thumbnail(img)

                # 压缩主图片（如果需要）
                processed_data, final_format = self._compress_image(img, img_format)

                # 转换为 base64
                base64_data = base64.b64encode(processed_data).decode("utf-8")
                base64_thumbnail = base64.b64encode(thumbnail_data).decode("utf-8")

                # 获取图片信息
                width, height = img.size
                mime_type = f"image/{final_format.lower()}"

                return {
                    "success": True,
                    "image_content": base64_data,
                    "thumbnail_content": base64_thumbnail,
                    "width": width,
                    "height": height,
                    "format": final_format,
                    "mime_type": mime_type,
                    "size_bytes": len(processed_data),
                    "original_filename": original_filename,
                }

        except Exception as e:
            logger.error(f"图片处理失败: {str(e)}")
            return {"success": False, "error": str(e)}

    def _validate_image_format(self, image_data: bytes) -> tuple[str, str]:
        """验证图片格式并返回格式信息"""
        try:
            with Image.open(io.BytesIO(image_data)) as img:
                return img.format, img.mode
        except Exception as e:
            raise ValueError(f"无效的图片格式: {str(e)}")

    def _fix_image_orientation(self, img: Image.Image) -> Image.Image:
        """根据EXIF信息修正图片方向"""
        try:
            if hasattr(img, "_getexif"):
                exif = img._getexif()
                if exif is not None:
                    for tag, value in exif.items():
                        if tag in ExifTags.TAGS and ExifTags.TAGS[tag] == "Orientation":
                            if value == 3:
                                img = img.rotate(180, expand=True)
                            elif value == 6:
                                img = img.rotate(270, expand=True)
                            elif value == 8:
                                img = img.rotate(90, expand=True)
                            break
        except Exception as e:
            logger.warning(f"修正图片方向失败，使用原始方向: {str(e)}")

        return img

    def _generate_thumbnail(self, img: Image.Image) -> bytes:
        """生成缩略图"""
        try:
            thumbnail = self._convert_to_rgb_for_export(img)

            # 生成缩略图，保持宽高比
            thumbnail.thumbnail(self.THUMBNAIL_SIZE, Image.Resampling.LANCZOS)

            # 转换为JPEG格式
            with io.BytesIO() as output:
                thumbnail.save(output, format="JPEG", quality=85, optimize=True)
                return output.getvalue()

        except Exception as e:
            logger.error(f"生成缩略图失败: {str(e)}")
            # 如果缩略图生成失败，返回一个1x1的透明图片
            with io.BytesIO() as output:
                empty_img = Image.new("RGB", (1, 1), color="white")
                empty_img.save(output, format="JPEG", quality=85)
                return output.getvalue()

    def _convert_to_rgb_for_export(self, img: Image.Image) -> Image.Image:
        """转换为 RGB，同时把透明像素按白底合成，避免隐藏颜色变成可见像素。"""
        if img.mode == "RGB":
            return img.copy()

        has_alpha = img.mode in ("RGBA", "LA") or (img.mode == "P" and "transparency" in img.info)
        if not has_alpha:
            return img.convert("RGB")

        rgba_img = img.convert("RGBA")
        background = Image.new("RGBA", rgba_img.size, (255, 255, 255, 255))
        background.alpha_composite(rgba_img)
        return background.convert("RGB")

    def _compress_image(self, img: Image.Image, original_format: str) -> tuple[bytes, str]:
        """
        压缩图片，如果超过大小限制

        Args:
            img: PIL Image对象
            original_format: 原始格式

        Returns:
            Tuple[bytes, str]: (压缩后的图片数据, 最终格式)
        """
        processed_img = self._convert_to_rgb_for_export(img)

        # 尝试保持原始格式，但优先使用JPEG（更好的压缩）
        target_format = "JPEG" if original_format != "PNG" else "PNG"

        # 初始质量设置
        quality = 85

        with io.BytesIO() as output:
            # 第一次保存以检查大小
            processed_img.save(output, format=target_format, quality=quality, optimize=True)
            compressed_data = output.getvalue()

            # 如果文件大小合适，直接返回
            if len(compressed_data) <= self.MAX_FILE_SIZE:
                return compressed_data, target_format

            # 如果文件太大，逐步降低质量
            while len(compressed_data) > self.MAX_FILE_SIZE and quality > 10:
                quality -= 10
                output.seek(0)
                output.truncate(0)
                processed_img.save(output, format=target_format, quality=quality, optimize=True)
                compressed_data = output.getvalue()

            # 如果质量降到最低仍然太大，尝试缩小尺寸
            if len(compressed_data) > self.MAX_FILE_SIZE:
                # 逐步缩小尺寸
                scale_factor = 0.9
                while len(compressed_data) > self.MAX_FILE_SIZE and scale_factor > 0.3:
                    new_width = int(processed_img.width * scale_factor)
                    new_height = int(processed_img.height * scale_factor)
                    resized_img = processed_img.resize((new_width, new_height), Image.Resampling.LANCZOS)

                    output.seek(0)
                    output.truncate(0)
                    resized_img.save(output, format=target_format, quality=85, optimize=True)
                    compressed_data = output.getvalue()

                    scale_factor -= 0.1

            return compressed_data, target_format


# 全局实例
image_processor = ImageProcessor()


def process_uploaded_image(image_data: bytes, filename: str = "") -> dict:
    """
    处理上传的图片（便捷函数）

    Args:
        image_data: 图片二进制数据
        filename: 文件名

    Returns:
        dict: 处理结果
    """
    return image_processor.process_image(image_data, filename)