from __future__ import annotations import json import os import uuid from dataclasses import dataclass from datetime import datetime, timezone from typing import Awaitable, Callable from urllib.parse import urlparse from app.config import settings from app.services.generation.pipeline.owner_service import GenerationOwner 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: GenerationOwner) -> 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 _normalize_image_extension(record: GenerationOwner) -> str: output_format = "" try: snapshot = json.loads(getattr(record, "engine_snapshot_json", None) or "{}") if isinstance(snapshot, dict): output_format = str(snapshot.get("output_format") or "").strip().lower() except Exception: output_format = "" if output_format in {"jpg", "jpeg"}: return "jpg" if output_format in {"png", "webp"}: return output_format remote_url = str(getattr(record, "remote_result_url", None) or "") try: suffix = os.path.splitext(urlparse(remote_url).path or "")[1].lower().lstrip(".") except Exception: suffix = "" if suffix in {"jpg", "jpeg"}: return "jpg" if suffix in {"png", "webp"}: return suffix return "jpg" 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, *, execution_guard: Callable[[], Awaitable[None]] | None = None, ) -> 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, execution_guard=execution_guard, ) if not _is_valid_file(part_path): raise RuntimeError("图片下载完成但临时文件为空") if execution_guard is not None: await execution_guard() 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, *, execution_guard: Callable[[], Awaitable[None]] | None = None, ) -> 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, execution_guard=execution_guard, ) if not _is_valid_file(part_path): raise RuntimeError("视频下载完成但临时文件为空") if execution_guard is not None: await execution_guard() os.replace(part_path, final_path) return final_path except Exception: _safe_remove(part_path) raise async def download_video_upscale_source( record: GenerationOwner, *, execution_guard: Callable[[], Awaitable[None]] | None = None, ) -> 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, execution_guard=execution_guard, ) if not _is_valid_file(part_path): raise RuntimeError("超分源视频下载完成但临时文件为空") await probe_video(part_path) if execution_guard is not None: await execution_guard() 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: GenerationOwner, *, execution_guard: Callable[[], Awaitable[None]] | None = None, ) -> 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) extension = _normalize_image_extension(record) dest = os.path.join(dest_dir, f"{record.id}.{extension}") 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, execution_guard=execution_guard, ) return DownloadedGenerationResult( url=f"/generate/images/{date_dir}/{record.id}.{extension}", 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, execution_guard=execution_guard, ) 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"生成资源视频封面 task_id={record.id}", ) if execution_guard is not None: await execution_guard() 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, )