xiaoxia-saas/apps/worker/video_processing/oss_helpers.py

"""OSS 工具函数 — 从 generation.py / edit_plan_generation.py 提取的共享 OSS 操作.

提供 OSS 配置读取、Bucket 创建、素材上传/下载、asset_id → 本地路径解析
等能力，供 render_edit_plan 和 generate_video 共同复用。
"""

from __future__ import annotations

import hashlib
import logging
import os
from pathlib import Path
from typing import Optional
from urllib.parse import urlparse

import oss2

logger = logging.getLogger(__name__)


# ── OSS 配置 ──────────────────────────────────────────────────────────────────


def oss_settings() -> tuple[str, str, str, str] | None:
    """获取 OSS 配置。

    Returns:
        (access_key_id, access_key_secret, endpoint, bucket_name) 元组，
        配置缺失时返回 None。
    """
    access_key_id = os.getenv("OSS_ACCESS_KEY_ID")
    access_key_secret = os.getenv("OSS_ACCESS_KEY_SECRET")
    endpoint = os.getenv("OSS_ENDPOINT")
    bucket_name = os.getenv("OSS_BUCKET_NAME")
    if not all([access_key_id, access_key_secret, endpoint, bucket_name]):
        return None
    return access_key_id, access_key_secret, endpoint, bucket_name


def oss_bucket() -> oss2.Bucket | None:
    """获取 OSS Bucket 实例。

    Returns:
        oss2.Bucket 实例，配置缺失时返回 None。
    """
    settings = oss_settings()
    if settings is None:
        return None
    access_key_id, access_key_secret, endpoint, bucket_name = settings
    return oss2.Bucket(oss2.Auth(access_key_id, access_key_secret), endpoint, bucket_name)


def normalize_storage_key(storage_key_or_url: str) -> str:
    """标准化存储键 — 如果是完整 URL 则提取 path 部分。

    Examples:
        "https://bucket.oss-cn-hangzhou.aliyuncs.com/path/to/file.mp4"
        → "path/to/file.mp4"
        "path/to/file.mp4" → "path/to/file.mp4"
    """
    if storage_key_or_url.startswith(("http://", "https://")):
        return urlparse(storage_key_or_url).path.lstrip("/")
    return storage_key_or_url.lstrip("/")


# ── 上传 / 下载 ───────────────────────────────────────────────────────────────


def download_asset(asset_storage_key: str, local_path: Path) -> bool:
    """从 OSS 下载素材文件到本地路径。

    Args:
        asset_storage_key: 素材的存储键（或完整 URL）
        local_path: 本地保存路径

    Returns:
        True 表示下载成功，False 表示失败。
    """
    bucket = oss_bucket()
    if bucket is None:
        return False
    try:
        bucket.get_object_to_file(normalize_storage_key(asset_storage_key), str(local_path))
        return local_path.exists() and local_path.stat().st_size > 0
    except Exception:
        logger.exception("下载素材失败: %s", asset_storage_key)
        return False


def upload_to_oss(local_path: Path, storage_key: str) -> str | None:
    """上传文件到 OSS，返回公开 URL。

    Args:
        local_path: 本地文件路径
        storage_key: 目标存储键

    Returns:
        公开访问 URL，上传失败或 OSS 未配置时返回 None。
    """
    bucket = oss_bucket()
    if bucket is None:
        return None
    try:
        bucket.put_object_from_file(storage_key, str(local_path))
        settings = oss_settings()
        if settings:
            _, _, endpoint, bucket_name = settings
            return f"https://{bucket_name}.{endpoint.replace('https://', '').replace('http://', '')}/{storage_key}"
        return None
    except Exception:
        logger.exception("上传 OSS 失败: %s", storage_key)
        return None


# ── Asset 解析 ────────────────────────────────────────────────────────────────


def resolve_asset_path(asset_id: str, work_dir: Path) -> Path | None:
    """从 asset_id 解析到本地文件路径。

    策略（按优先级）：
    1. 如果 asset_id 是本地绝对路径（/var/storage/...）→ 直接返回
    2. 如果 work_dir 下已有缓存文件 → 返回缓存路径
    3. 从 OSS 下载到 work_dir/{hash}.mp4 → 返回下载路径
    4. 下载失败 → 返回 None

    缓存策略：以 asset_id 的 SHA256 前 16 位为文件名，避免重复下载。
    """
    # 1. 本地绝对路径
    if asset_id.startswith("/") and os.path.exists(asset_id):
        return Path(asset_id)

    # 2. 缓存命中
    cache_hash = hashlib.sha256(asset_id.encode()).hexdigest()[:16]
    cached_path = work_dir / f"{cache_hash}.mp4"
    if cached_path.exists() and cached_path.stat().st_size > 0:
        return cached_path

    # 3. 从 OSS 下载
    if download_asset(asset_id, cached_path):
        return cached_path

    return None


def resolve_asset_ids_to_paths(
    asset_ids: list[str],
    work_dir: Path,
) -> dict[str, Path]:
    """批量解析 asset_id → 本地路径。

    Args:
        asset_ids: 素材 ID 列表
        work_dir: 工作目录

    Returns:
        {asset_id: local_path} 映射，仅包含成功解析的条目。
    """
    result: dict[str, Path] = {}
    for aid in asset_ids:
        local_path = resolve_asset_path(aid, work_dir)
        if local_path:
            result[aid] = local_path
    return result