ozon-seller-kit/server/services/image_assets.py

"""商品图片生命周期：原图、当前水印版与存储回收的唯一入口。

每个图片素材只维护两个逻辑版本：original（无水印）与 current（页面展示/导出）。
未加水印时两者指向同一个对象；重加水印会先上传新 current，再切换引用并删除旧 current，
不保留历史水印版。
"""
from __future__ import annotations

import mimetypes
import uuid
from dataclasses import dataclass

from api.proxy import guess_referer
from config import get_settings
from models import ProductAsset
from services.storage import (
    download_bytes,
    get_storage,
    image_extension,
    local_path,
    normalize_image_bytes,
    read_local_bytes,
    storage_key_from_url,
)
from services.watermark import apply_watermark


@dataclass
class WatermarkUpdate:
    """水印切换后的旧对象；数据库提交成功后再回收。"""

    obsolete_key: str | None = None


class ImageAssetService:
    """将 ProductAsset 与七牛/本地对象存储的生命周期收敛到同一服务。"""

    def __init__(self) -> None:
        self.storage = get_storage()

    @staticmethod
    def asset_prefix(asset: ProductAsset) -> str:
        return f"products/{asset.product_id}/assets/{asset.id}/"

    @staticmethod
    def product_prefix(product_id: str) -> str:
        return f"products/{product_id}/"

    def _key(self, asset: ProductAsset, name: str, data: bytes, content_type: str) -> str:
        return f"{self.asset_prefix(asset)}{name}{image_extension(data, content_type)}"

    async def _read(self, url: str) -> tuple[bytes, str]:
        path = local_path(url)
        if path is not None:
            return read_local_bytes(path), mimetypes.guess_type(path.name)[0] or "image/jpeg"
        return await download_bytes(url, referer=guess_referer(url))

    async def save_original(
        self,
        asset: ProductAsset,
        data: bytes,
        content_type: str = "",
        max_bytes: int | None = None,
    ) -> str:
        """写入该素材唯一的无水印原图，并将 current 同步到原图。"""
        if asset.type == "img":
            normalized = normalize_image_bytes(data, max_bytes=max_bytes)
            data = normalized.data
            content_type = normalized.content_type
            asset.width = normalized.width
            asset.height = normalized.height
        key = self._key(asset, "original", data, content_type)
        url = await self.storage.save_bytes(data, key, content_type)
        asset.original_url = url
        asset.original_storage_key = key
        asset.stored_url = url
        asset.current_storage_key = key
        asset.watermark_enabled = False
        asset.watermark_config = None
        return url

    async def save_original_from_url(self, asset: ProductAsset, url: str) -> str:
        data, content_type = await self._read(url)
        return await self.save_original(
            asset,
            data,
            content_type,
            max_bytes=get_settings().image_download_max_bytes,
        )

    async def ensure_original(self, asset: ProductAsset) -> tuple[bytes, str]:
        """为旧记录补齐原图字段；首次触碰旧素材时迁入新的商品/素材路径。"""
        original_url = asset.original_url or asset.stored_url or asset.source_url
        if not original_url:
            raise ValueError("图片原图地址为空")

        # 新格式已在当前素材路径中，无需复制。
        if asset.original_storage_key and asset.original_url:
            return await self._read(asset.original_url)

        data, content_type = await self._read(original_url)
        await self.save_original(asset, data, content_type)
        return data, content_type

    async def apply_watermark(self, asset: ProductAsset, opts: dict) -> WatermarkUpdate:
        """从 original 合成新水印版，调用方提交数据库后再删除返回的旧水印对象。"""
        previous_current_key = asset.current_storage_key or storage_key_from_url(asset.stored_url or "")
        data, content_type = await self.ensure_original(asset)
        watermarked = apply_watermark(data, opts)
        key = self._key(asset, f"watermark-{uuid.uuid4().hex}", watermarked, content_type)
        url = await self.storage.save_bytes(watermarked, key, content_type)

        obsolete_key = previous_current_key
        # 未加水印时 current 就是 original，绝不能作为旧水印删除。
        if obsolete_key == asset.original_storage_key:
            obsolete_key = None

        asset.stored_url = url
        asset.current_storage_key = key
        asset.watermark_enabled = True
        asset.watermark_config = opts
        return WatermarkUpdate(obsolete_key=obsolete_key)

    async def adopt_generated_versions(
        self,
        asset: ProductAsset,
        original_url: str,
        current_url: str,
        watermark_config: dict | None = None,
    ) -> None:
        """将任务临时输出迁入正式的商品/素材路径，不让 suites 临时对象长期残留。"""
        await self.save_original_from_url(asset, original_url)
        if current_url == original_url:
            return
        data, content_type = await self._read(current_url)
        key = self._key(asset, f"watermark-{uuid.uuid4().hex}", data, content_type)
        asset.stored_url = await self.storage.save_bytes(data, key, content_type)
        asset.current_storage_key = key
        asset.watermark_enabled = True
        asset.watermark_config = watermark_config or {}

    async def remove_watermark(self, asset: ProductAsset) -> WatermarkUpdate:
        """切回无水印原图，调用方提交数据库后再删除旧水印对象。"""
        previous_current_key = asset.current_storage_key or storage_key_from_url(asset.stored_url or "")
        await self.ensure_original(asset)
        obsolete_key = previous_current_key
        asset.stored_url = asset.original_url
        asset.current_storage_key = asset.original_storage_key
        asset.watermark_enabled = False
        asset.watermark_config = None
        if obsolete_key == asset.original_storage_key:
            obsolete_key = None
        return WatermarkUpdate(obsolete_key=obsolete_key)

    async def delete_key(self, key: str | None) -> None:
        if key:
            await self.storage.delete(key)

    async def delete_temporary_urls(self, urls: list[str]) -> None:
        """只回收 suites/ 临时输出，绝不删除作为参考图传入的正式商品素材。"""
        for url in set(urls):
            key = storage_key_from_url(url)
            if key and key.startswith("suites/"):
                await self.delete_key(key)

    async def delete_asset_media(self, asset: ProductAsset) -> int:
        """删除单个素材的当前格式目录，并兼容清理由旧数据留下的路径外对象。"""
        explicit_keys = {
            key for key in (
                asset.original_storage_key,
                asset.current_storage_key,
                storage_key_from_url(asset.original_url or ""),
                storage_key_from_url(asset.stored_url or ""),
            ) if key
        }
        deleted = await self.storage.delete_prefix(self.asset_prefix(asset))
        for key in explicit_keys:
            if not key.startswith(self.asset_prefix(asset)):
                await self.delete_key(key)
                deleted += 1
        return deleted

    async def delete_product_media(self, product_id: str, assets: list[ProductAsset]) -> int:
        """删除商品下所有素材；先清理旧 key，再按商品前缀兜底清理孤儿对象。"""
        deleted = 0
        for asset in assets:
            deleted += await self.delete_asset_media(asset)
        deleted += await self.storage.delete_prefix(self.product_prefix(product_id))
        return deleted


def get_image_asset_service() -> ImageAssetService:
    return ImageAssetService()