video-gen/video-gen-api/app/services/generation/download_service.py

from __future__ import annotations

import os
import uuid
from dataclasses import dataclass
from datetime import datetime, timezone

from app.config import settings
from app.models.chat_generation_task import ChatGenerationTask
from app.services.image_gen import download_image
from app.services.provider_limit import provider_limit
from app.services.resource_accounting_service import safe_file_size
from app.services.video_cover_service import create_video_cover_for_local_video
from app.services.video_gen import download_video
from app.services.video_upscale.media_service import build_part_mp4_path, probe_video


@dataclass(slots=True)
class DownloadedGenerationResult:
    url: str
    storage_path: str | None
    file_size_bytes: int
    resource_type: str
    storage_type: str = "local"
    cover_url: str | None = None
    cover_storage_path: str | None = None


def _to_aware_utc(value: datetime | None) -> datetime | None:
    if value is None:
        return None
    if value.tzinfo is None:
        return value.replace(tzinfo=timezone.utc)
    return value.astimezone(timezone.utc)


def _build_storage_date_dir(record: ChatGenerationTask) -> str:
    fixed = (getattr(record, "download_storage_date_dir", None) or "").strip().strip("/")
    if fixed:
        return fixed
    created_at = _to_aware_utc(getattr(record, "created_at", None)) or datetime.now(timezone.utc)
    return created_at.strftime("%Y/%m/%d")


def _make_part_path(final_path: str) -> str:
    return f"{final_path}.{uuid.uuid4().hex}.part"


def _is_valid_file(path: str | None) -> bool:
    if not path:
        return False
    try:
        return os.path.isfile(path) and os.path.getsize(path) > 0
    except OSError:
        return False


def _safe_remove(path: str | None) -> None:
    if not path:
        return
    try:
        if os.path.exists(path):
            os.remove(path)
    except OSError:
        pass


async def _download_image_atomically(remote_url: str, final_path: str) -> str:
    if _is_valid_file(final_path):
        return final_path

    os.makedirs(os.path.dirname(final_path), exist_ok=True)
    part_path = _make_part_path(final_path)
    try:
        await download_image(remote_url, part_path)
        if not _is_valid_file(part_path):
            raise RuntimeError("图片下载完成但临时文件为空")
        os.replace(part_path, final_path)
        return final_path
    except Exception:
        _safe_remove(part_path)
        raise


async def _download_video_atomically(remote_url: str, final_path: str) -> str:
    if _is_valid_file(final_path):
        return final_path

    os.makedirs(os.path.dirname(final_path), exist_ok=True)
    part_path = _make_part_path(final_path)
    try:
        await download_video(remote_url, part_path)
        if not _is_valid_file(part_path):
            raise RuntimeError("视频下载完成但临时文件为空")
        os.replace(part_path, final_path)
        return final_path
    except Exception:
        _safe_remove(part_path)
        raise




async def download_video_upscale_source(record: ChatGenerationTask) -> DownloadedGenerationResult:
    """下载超分源视频。

    源视频只供后处理使用，不生成封面，也不作为用户 GeneratedResource。
    临时文件保持 .part.mp4 后缀，校验成功后原子重命名为 .source.mp4。
    """
    if not record.remote_result_url:
        raise ValueError("缺少远程结果URL")

    date_dir = _build_storage_date_dir(record)
    dest_dir = os.path.join(settings.STORAGE_LOCAL_PATH, "_upscale_source", date_dir)
    os.makedirs(dest_dir, exist_ok=True)
    dest = os.path.join(dest_dir, f"{record.id}.source.mp4")

    if not _is_valid_file(dest):
        part_path = build_part_mp4_path(dest)
        try:
            async with provider_limit("result_download", settings.RESULT_DOWNLOAD_MAX_CONCURRENCY):
                await download_video(record.remote_result_url, part_path)
            if not _is_valid_file(part_path):
                raise RuntimeError("超分源视频下载完成但临时文件为空")
            await probe_video(part_path)
            os.replace(part_path, dest)
        except Exception:
            _safe_remove(part_path)
            raise
    else:
        await probe_video(dest)

    return DownloadedGenerationResult(
        url=f"/generate/videos/_upscale_source/{date_dir}/{record.id}.source.mp4",
        storage_path=dest,
        file_size_bytes=safe_file_size(dest),
        resource_type="video",
    )


async def download_generation_result(record: ChatGenerationTask) -> DownloadedGenerationResult:
    if not record.remote_result_url:
        raise ValueError("缺少远程结果URL")

    date_dir = _build_storage_date_dir(record)

    if record.gen_type == "image":
        dest_dir = os.path.join(settings.STORAGE_IMAGE_LOCAL_PATH, date_dir)
        os.makedirs(dest_dir, exist_ok=True)
        dest = os.path.join(dest_dir, f"{record.id}.png")

        async with provider_limit("result_download", settings.RESULT_DOWNLOAD_MAX_CONCURRENCY):
            await _download_image_atomically(record.remote_result_url if record.remote_result_url else "", dest)

        return DownloadedGenerationResult(
            url=f"/generate/images/{date_dir}/{record.id}.png",
            storage_path=dest,
            file_size_bytes=safe_file_size(dest),
            resource_type="image",
        )

    dest_dir = os.path.join(settings.STORAGE_LOCAL_PATH, date_dir)
    os.makedirs(dest_dir, exist_ok=True)
    dest = os.path.join(dest_dir, f"{record.id}.mp4")

    async with provider_limit("result_download", settings.RESULT_DOWNLOAD_MAX_CONCURRENCY):
        await _download_video_atomically(record.remote_result_url if record.remote_result_url else "", dest)

    cover_url, cover_storage_path = create_video_cover_for_local_video(
        record_id=record.id,
        video_path=dest,
        date_dir=date_dir,
        log_prefix=f"ChatGenerationTask视频封面生成 task_id={record.id}",
    )

    return DownloadedGenerationResult(
        url=f"/generate/videos/{date_dir}/{record.id}.mp4",
        storage_path=dest,
        file_size_bytes=safe_file_size(dest),
        resource_type="video",
        cover_url=cover_url,
        cover_storage_path=cover_storage_path,
    )