src/preprocessing/image_processor.py

"""
ImageProcessor — 图片预处理。

加载、验证、缩放图片, 并支持转为 base64 编码。
"""

import base64
import logging
import uuid
from io import BytesIO
from pathlib import Path
from typing import Dict, List, Optional, Tuple

from PIL import Image

from config.settings import settings

logger = logging.getLogger(__name__)

# 支持的图片格式
SUPPORTED_IMAGE_FORMATS = {
    ".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp", ".tiff", ".tif",
}


class ImageProcessor:
    """图片预处理器。

    功能:
      - 加载并验证图片
      - 等比缩放 (保持宽高比)
      - 转为 JPEG 字节或 base64 字符串
      - 生成用于媒体展示的缩略图

    Usage:
        proc = ImageProcessor(target_size=256)
        result = proc.process("photo.jpg")  # -> {"image_bytes": ..., "base64": ..., "metadata": ...}
    """

    def __init__(
        self,
        target_size: int = 256,
        quality: int = 85,
        thumbnail_dim: int = 256,
    ):
        self.target_size = target_size
        self.quality = quality
        self.thumbnail_dim = thumbnail_dim

    def process(self, image_path: str) -> dict:
        """处理单张图片。

        Args:
            image_path: 图片文件路径。

        Returns:
            包含以下字段的字典:
                - chunk_id: 唯一 ID
                - image_bytes: JPEG 编码的图片字节
                - base64: base64 字符串
                - metadata: 元数据 (尺寸、格式、来源等)
                - thumbnail_base64: 缩略图 base64
                - modality: "image"

        Raises:
            FileNotFoundError: 文件不存在。
            ValueError: 文件格式不支持或内容无效。
        """
        path = Path(image_path)
        if not path.exists():
            raise FileNotFoundError(f"图片文件不存在: {image_path}")

        suffix = path.suffix.lower()
        if suffix not in SUPPORTED_IMAGE_FORMATS:
            raise ValueError(
                f"不支持的图片格式: {suffix}。支持的格式: {SUPPORTED_IMAGE_FORMATS}"
            )

        # 预检文件头魔数，防止非图片文件以图片扩展名上传
        self._validate_image_bytes(path)

        try:
            img = Image.open(image_path)
            img.load()  # 强制加载像素数据，触发真实解码
        except Exception as e:
            raise ValueError(
                f"无法解析图片文件，文件可能已损坏或格式与扩展名不匹配: {path.name}"
            ) from e

        original_size = img.size

        # 转 RGB
        if img.mode in ("RGBA", "P", "LA", "PA"):
            img = img.convert("RGB")

        # 缩放到目标尺寸 (居中裁剪)
        img_resized = self._resize_to_target(img, self.target_size)

        # 缩略图
        thumbnail = self._resize_to_target(img.copy(), self.thumbnail_dim)

        # 编码为 JPEG
        main_bytes = self._encode_jpeg(img_resized)
        thumbnail_bytes = self._encode_jpeg(thumbnail, quality=75)

        chunk_id = str(uuid.uuid4())

        return {
            "chunk_id": chunk_id,
            "image_bytes": main_bytes,
            "base64": base64.b64encode(main_bytes).decode("utf-8"),
            "thumbnail_base64": base64.b64encode(thumbnail_bytes).decode("utf-8"),
            "metadata": {
                "chunk_id": chunk_id,
                "source_file": str(path.resolve()),
                "source_file_name": path.name,
                "original_size": original_size,
                "resized_size": (self.target_size, self.target_size),
                "format": suffix.lstrip("."),
                "modality": "image",
                "content_type": "image",
            },
            "modality": "image",
        }

    def process_batch(self, image_paths: List[str]) -> List[dict]:
        """批量处理图片。"""
        results = []
        for p in image_paths:
            try:
                results.append(self.process(p))
            except Exception as e:
                logger.error(f"处理图片失败 {p}: {e}")
        return results

    def _resize_to_target(self, img: Image.Image, size: int) -> Image.Image:
        """居中裁剪并缩放到指定尺寸 (size x size)。"""
        w, h = img.size

        # 等比缩放使短边 = size
        if w < h:
            new_w = size
            new_h = int(h * (size / w))
        else:
            new_h = size
            new_w = int(w * (size / h))
        img = img.resize((new_w, new_h), Image.LANCZOS)

        # 居中裁剪
        left = (new_w - size) // 2
        top = (new_h - size) // 2
        return img.crop((left, top, left + size, top + size))

    def _encode_jpeg(self, img: Image.Image, quality: Optional[int] = None) -> bytes:
        """编码为 JPEG 字节。"""
        quality = quality or self.quality
        buf = BytesIO()
        img.save(buf, format="JPEG", quality=quality)
        return buf.getvalue()

    @staticmethod
    def _validate_image_bytes(path: Path):
        """通过文件头魔数验证是否为有效图片。"""
        MAGIC_BYTES = {
            b"\xff\xd8\xff": "JPEG",
            b"\x89PNG\r\n\x1a\n": "PNG",
            b"GIF87a": "GIF",
            b"GIF89a": "GIF",
            b"RIFF": "WEBP",
            b"BM": "BMP",
            b"II*\x00": "TIFF",
            b"MM\x00*": "TIFF",
        }

        try:
            with open(path, "rb") as f:
                header = f.read(12)
        except OSError as e:
            raise ValueError(f"无法读取文件: {e}")

        if len(header) < 4:
            raise ValueError(f"文件过小，不是有效图片: {path.name}")

        if header[:4] == b"RIFF" and header[8:12] == b"WEBP":
            return

        for magic, fmt in MAGIC_BYTES.items():
            if header[:len(magic)] == magic:
                return

        raise ValueError(
            f"文件头不匹配任何已知图片格式，文件可能不是有效图片: {path.name}"
        )

    @staticmethod
    def get_image_base64(image_path: str, target_size: int = 256) -> str:
        """快速获取图片 base64 (不生成缩略图)。"""
        proc = ImageProcessor(target_size=target_size)
        result = proc.process(image_path)
        return result["base64"]

    @staticmethod
    def bytes_to_base64(image_bytes: bytes, target_size: int = 256) -> str:
        """将图片 bytes 转为 base64。"""
        proc = ImageProcessor(target_size=target_size)
        img = Image.open(BytesIO(image_bytes))
        if img.mode in ("RGBA", "P", "LA", "PA"):
            img = img.convert("RGB")
        img_resized = proc._resize_to_target(img, target_size)
        jpeg_bytes = proc._encode_jpeg(img_resized)
        return base64.b64encode(jpeg_bytes).decode("utf-8")