video-gen/video-gen-api/app/services/home_material/storage.py

from __future__ import annotations

import os
import shutil
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Iterable

from fastapi import HTTPException, UploadFile, status

from app.config import settings
from app.enums.home_material import (
    HOME_MATERIAL_IMAGE_EXTENSIONS,
    HOME_MATERIAL_VIDEO_EXTENSIONS,
    HOME_MATERIAL_WATERMARK_EXTENSIONS,
    HomeMaterialMediaType,
)
from app.utils.id_gen import generate_id


@dataclass(frozen=True)
class StoredFile:
    storage_path: str
    file_url: str
    file_name: str
    file_size_bytes: int
    suffix: str


class HomeMaterialStorageService:
    """首页素材文件存储服务。只处理文件、目录、URL，不访问数据库。"""

    def __init__(self) -> None:
        self.upload_root = Path(settings.UPLOAD_LOCAL_PATH).resolve()
        self.home_root = self.upload_root / "home_materials"

    def _date_dir(self) -> str:
        return datetime.now().strftime("%Y/%m/%d")

    def _safe_suffix(self, filename: str | None) -> str:
        return Path(filename or "").suffix.lower()

    def _assert_extension(self, suffix: str, allowed: Iterable[str], label: str) -> None:
        if suffix not in set(allowed):
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"{label}文件类型不支持：{suffix or '无扩展名'}",
            )

    def _limit_mb(self, media_type: HomeMaterialMediaType | str) -> int:
        if media_type == HomeMaterialMediaType.VIDEO or str(media_type) == HomeMaterialMediaType.VIDEO.value:
            return int(getattr(settings, "HOME_MATERIAL_MAX_VIDEO_MB", 300))
        return int(getattr(settings, "HOME_MATERIAL_MAX_IMAGE_MB", 20))

    def _build_target(self, relative_dir: str, suffix: str) -> tuple[Path, str]:
        file_id = generate_id()
        relative = Path(relative_dir) / self._date_dir() / f"{file_id}{suffix}"
        path = self.home_root / relative
        path.parent.mkdir(parents=True, exist_ok=True)
        return path, f"/uploads/home_materials/{relative.as_posix()}"

    async def save_upload_file(
        self,
        file: UploadFile,
        *,
        relative_dir: str,
        allowed_extensions: Iterable[str],
        label: str,
        max_mb: int | None = None,
    ) -> StoredFile:
        suffix = self._safe_suffix(file.filename)
        self._assert_extension(suffix, allowed_extensions, label)
        target_path, file_url = self._build_target(relative_dir, suffix)
        max_bytes = int(max_mb or 0) * 1024 * 1024 if max_mb else None
        tmp_path = target_path.with_name(target_path.name + ".part")
        size = 0
        try:
            with tmp_path.open("wb") as out:
                while True:
                    chunk = await file.read(1024 * 1024)
                    if not chunk:
                        break
                    size += len(chunk)
                    if max_bytes and size > max_bytes:
                        raise HTTPException(
                            status_code=status.HTTP_400_BAD_REQUEST,
                            detail=f"{label}文件大小超限，最大 {max_mb}MB",
                        )
                    out.write(chunk)
            os.replace(tmp_path, target_path)
        except Exception:
            try:
                if tmp_path.exists():
                    tmp_path.unlink()
            finally:
                raise
        finally:
            await file.seek(0)

        return StoredFile(
            storage_path=str(target_path),
            file_url=file_url,
            file_name=file.filename or target_path.name,
            file_size_bytes=size,
            suffix=suffix,
        )

    async def save_asset_file(self, file: UploadFile, media_type: HomeMaterialMediaType | str) -> StoredFile:
        media = HomeMaterialMediaType(media_type)
        if media == HomeMaterialMediaType.IMAGE:
            return await self.save_upload_file(
                file,
                relative_dir="original/images",
                allowed_extensions=HOME_MATERIAL_IMAGE_EXTENSIONS,
                label="图片素材",
                max_mb=self._limit_mb(media),
            )
        return await self.save_upload_file(
            file,
            relative_dir="original/videos",
            allowed_extensions=HOME_MATERIAL_VIDEO_EXTENSIONS,
            label="视频素材",
            max_mb=self._limit_mb(media),
        )

    async def save_watermark_file(self, file: UploadFile) -> StoredFile:
        return await self.save_upload_file(
            file,
            relative_dir="watermarks",
            allowed_extensions=HOME_MATERIAL_WATERMARK_EXTENSIONS,
            label="水印图片",
            max_mb=int(getattr(settings, "HOME_MATERIAL_MAX_WATERMARK_MB", 10)),
        )

    def build_watermarked_target(self, asset_id: str, media_type: HomeMaterialMediaType | str) -> tuple[str, str]:
        media = HomeMaterialMediaType(media_type)
        suffix = ".png" if media == HomeMaterialMediaType.IMAGE else ".mp4"
        folder = "watermarked/images" if media == HomeMaterialMediaType.IMAGE else "watermarked/videos"
        relative = Path(folder) / self._date_dir() / f"{asset_id}{suffix}"
        path = self.home_root / relative
        path.parent.mkdir(parents=True, exist_ok=True)
        return str(path), f"/uploads/home_materials/{relative.as_posix()}"

    def build_cover_target(self, asset_id: str) -> tuple[str, str]:
        relative = Path("covers") / self._date_dir() / f"{asset_id}.jpg"
        path = self.home_root / relative
        path.parent.mkdir(parents=True, exist_ok=True)
        return str(path), f"/uploads/home_materials/{relative.as_posix()}"

    @staticmethod
    def file_size(path: str | None) -> int | None:
        if not path:
            return None
        try:
            return Path(path).stat().st_size
        except FileNotFoundError:
            return None

    @staticmethod
    def safe_remove(path: str | None) -> None:
        if not path:
            return
        try:
            Path(path).unlink(missing_ok=True)
        except Exception:
            pass


storage_service = HomeMaterialStorageService()