xiaoxia-saas/apps/api/app/api/routes/generation_tasks.py

import logging
import random
import uuid
from typing import Any

from app.api.routes._helpers import check_project_access
from app.auth import AuthenticatedUser, get_current_user
from app.core.storage import OSSStorageService, get_storage_service
from app.core.task_enqueue import (
    GLOBAL_PENDING_LIMIT,
    USER_PENDING_LIMIT,
    GlobalQueueFull,
    UserPendingLimitExceeded,
    safe_enqueue_generation_task,
)
from app.dependencies import (
    get_asset_library_repository,
    get_asset_repository,
    get_generated_video_repository,
    get_generation_task_repository,
    get_project_repository,
)
from app.schemas.generated_video import (
    GeneratedVideoResponse,
    ListGeneratedVideosResponse,
)
from app.schemas.generation_task import (
    BatchGenerationTaskResponse,
    CreateGenerationTaskRequest,
    GenerationTaskResponse,
    ListGenerationTasksResponse,
)
from fastapi import APIRouter, Depends, HTTPException

from packages.application import (
    CreateGenerationTaskCommand,
    CreateGenerationTaskUseCase,
    GetGenerationTaskUseCase,
    ListGeneratedVideosByTaskUseCase,
)

logger = logging.getLogger(__name__)

router = APIRouter()


def _to_generation_task_response(task) -> GenerationTaskResponse:
    return GenerationTaskResponse(
        id=task.id,
        project_id=task.project_id,
        asset_library_id=task.asset_library_id,
        strategy_id=task.strategy_id,
        voice_library_id=task.voice_library_id,
        template_id=task.template_id,
        asset_ids=task.asset_ids,
        title_ids=task.title_ids,
        voice_ids=task.voice_ids,
        source_edit_plan_id=task.source_edit_plan_id or "",
        asset_select_mode=getattr(task, "asset_select_mode", ""),
        batch_id=getattr(task, "batch_id", ""),
        video_title=getattr(task, "video_title", ""),
        resolution=getattr(task, "resolution", ""),
        bgm_config=getattr(task, "bgm_config", {}) or {},
        logs=getattr(task, "logs", "[]"),
        status=task.status,
        progress=task.progress,
        result_count=task.result_count,
        error_message=task.error_message,
    )


def _to_generated_video_response(item, download_url: str | None = None) -> GeneratedVideoResponse:
    return GeneratedVideoResponse(
        id=item.id,
        project_id=item.project_id,
        generation_task_id=item.generation_task_id,
        name=item.name,
        file_url=item.file_url,
        file_size=item.file_size,
        duration=item.duration,
        thumbnail_url=item.thumbnail_url,
        width=item.width,
        height=item.height,
        fps=item.fps,
        download_url=download_url,
    )


def _ensure_library_has_ready_video_assets(assets) -> None:
    ready_video_assets = [
        asset for asset in assets if asset.status.value == "ready" and asset.mime_type.startswith("video")
    ]
    if not ready_video_assets:
        raise HTTPException(
            status_code=422,
            detail="当前素材库没有 ready 状态的视频素材，请先上传并等待导入完成后再生成。",
        )


def _select_assets_from_library(
    assets: list,
    mode: str,
    count: int,
) -> list[str]:
    """根据选取模式从素材库中选取 ready 状态的视频素材 ID。

    Args:
        assets: 素材库中所有素材（Asset 实体列表）
        mode: 选取模式 — all=全部, random=随机, smart=智能匹配（多维度评分+多样性）
        count: 选取数量，0 表示全部（仅 random/smart 模式有效）

    Returns:
        选中的素材 ID 列表
    """
    ready_video_assets = [a for a in assets if a.status.value == "ready" and a.mime_type.startswith("video")]

    if not ready_video_assets:
        return []

    if mode == "random":
        selected = (
            ready_video_assets if count <= 0 else random.sample(ready_video_assets, min(count, len(ready_video_assets)))
        )
        return [a.id for a in selected]

    if mode == "smart":
        # 智能匹配：按质量分降序 + 时长降序作为tiebreaker
        # 注意：这里使用简单的 quality_score 排序保持向后兼容
        # 更复杂的4维评分+多样性策略由 SmartAssetSelector 服务提供（用于 AI 精选等场景）
        scored_assets = sorted(
            ready_video_assets,
            key=lambda a: (
                -(a.quality_score if a.quality_score is not None else 0.0),
                -(getattr(a, "duration", 0.0) or 0.0),
            ),
        )
        if count > 0:
            scored_assets = scored_assets[:count]
        return [a.id for a in scored_assets]

    # 默认 all 模式：返回全部 ready 视频素材
    return [a.id for a in ready_video_assets]


def _resolve_project_and_library(
    request: CreateGenerationTaskRequest,
    project_repository: Any,
    asset_library_repository: Any,
    asset_repository: Any,
    authenticated_user: AuthenticatedUser,
) -> tuple[str, str]:
    """解析 project_id 和 asset_library_id。

    支持两种模式：
    - 显式传入（向后兼容）
    - 从 asset_ids 反查 asset_library（模板模式）
    返回 (project_id, asset_library_id)。
    """
    project_id = request.project_id.strip()
    asset_library_id = request.asset_library_id.strip()

    # 模板模式：project_id 未提供时，从 asset_ids 反查所属 project
    if not project_id and request.asset_ids:
        first_asset_id = request.asset_ids[0]
        asset = asset_repository.find_by_id(first_asset_id)
        if asset is not None:
            project_id = asset.project_id
            if not asset_library_id:
                asset_library_id = asset.library_id

    # 向后兼容校验：project_id 已提供时验证权限
    if project_id:
        project = project_repository.find_by_id(project_id)
        if project is None:
            raise HTTPException(status_code=404, detail=f"Project {project_id} not found")
        if not project.can_access(authenticated_user.user.id):
            raise HTTPException(status_code=403, detail="Access denied to project")

    return project_id, asset_library_id


@router.post("/tasks", response_model=BatchGenerationTaskResponse)
def create_generation_task(
    request: CreateGenerationTaskRequest,
    authenticated_user: AuthenticatedUser = Depends(get_current_user),
    generation_task_repository: Any = Depends(get_generation_task_repository),
    project_repository: Any = Depends(get_project_repository),
    asset_library_repository: Any = Depends(get_asset_library_repository),
    asset_repository: Any = Depends(get_asset_repository),
) -> BatchGenerationTaskResponse:
    logger.info(
        "[生成任务] 接收请求: user_id=%s, template_id=%s, asset_count=%d, mode=%s, count=%d",
        authenticated_user.user.id,
        request.template_id,
        len(request.asset_ids),
        request.asset_select_mode,
        request.count,
    )

    try:
        project_id, asset_library_id = _resolve_project_and_library(
            request, project_repository, asset_library_repository, asset_repository, authenticated_user
        )
    except HTTPException as e:
        logger.warning("[生成任务] 校验失败: %s", e.detail)
        raise

    # asset_library 存在性校验（仅在提供了 asset_library_id 时）
    resolved_asset_ids: list[str] = list(request.asset_ids)
    if asset_library_id:
        library = asset_library_repository.get(asset_library_id)
        if library is None or (project_id and library.project_id != project_id):
            logger.warning("[生成任务] 素材库不存在: library_id=%s", asset_library_id)
            raise HTTPException(status_code=404, detail=f"AssetLibrary {asset_library_id} not found")

        assets = asset_repository.find_by_library(asset_library_id)
        try:
            _ensure_library_has_ready_video_assets(assets)
        except HTTPException as e:
            logger.warning("[生成任务] 素材校验失败: %s", e.detail)
            raise

        # 素材库自动匹配：当未显式指定 asset_ids 时，按模式自动选取
        if not resolved_asset_ids:
            resolved_asset_ids = _select_assets_from_library(
                assets,
                mode=request.asset_select_mode,
                count=request.asset_select_count,
            )
    elif project_id and not resolved_asset_ids and request.asset_select_mode in ("random", "smart"):
        # 项目级模式：未指定 asset_ids 且选择了 random/smart 模式时，也自动选取
        assets = asset_repository.find_by_project(project_id)
        if assets:
            resolved_asset_ids = _select_assets_from_library(
                assets,
                mode=request.asset_select_mode,
                count=request.asset_select_count,
            )
            if not resolved_asset_ids:
                raise HTTPException(
                    status_code=422,
                    detail="当前项目没有符合条件的视频素材，请先上传并等待导入完成后再生成。",
                )

    use_case = CreateGenerationTaskUseCase(generation_task_repository)
    count = request.count
    created_tasks = []
    failed_tasks = []
    user_id = authenticated_user.user.id
    # 同批次任务共享 batch_id，用于视频查重时批次内比对
    batch_id = uuid.uuid4().hex if count > 1 else ""

    # 预检查：批量提交前先看会不会超限，避免建一半才拒
    try:
        user_pending = generation_task_repository.count_pending_by_user(user_id)
        global_pending = generation_task_repository.count_pending_total()
        if user_pending + count > USER_PENDING_LIMIT:
            raise UserPendingLimitExceeded(
                user_id=user_id, pending_count=user_pending + count, limit=USER_PENDING_LIMIT
            )
        if global_pending + count > GLOBAL_PENDING_LIMIT:
            raise GlobalQueueFull(pending_count=global_pending + count, limit=GLOBAL_PENDING_LIMIT)
    except UserPendingLimitExceeded as e:
        raise HTTPException(
            status_code=429,
            detail=f"您的待处理任务过多（当前 {e.pending_count - count}/{e.limit}，本次提交 {count} 个），请等待完成后再提交",
        ) from e
    except GlobalQueueFull as e:
        raise HTTPException(
            status_code=503,
            detail="系统繁忙，请稍后再试",
        ) from e

    try:
        for _ in range(count):
            task = use_case.execute(
                CreateGenerationTaskCommand(
                    project_id=project_id,
                    asset_library_id=asset_library_id,
                    strategy_id=request.strategy_id,
                    voice_library_id=request.voice_library_id,
                    template_id=request.template_id,
                    asset_ids=resolved_asset_ids,
                    title_ids=request.title_ids,
                    voice_ids=request.voice_ids,
                    created_by_user_id=user_id,
                    source_edit_plan_id=request.source_edit_plan_id,
                    asset_select_mode=request.asset_select_mode,
                    batch_id=batch_id,
                    video_title=request.video_title,
                    resolution=request.resolution,
                    bgm_config=request.bgm_config,
                    auto_retry_enabled=request.auto_retry_enabled,
                    auto_retry_max=request.auto_retry_max,
                )
            )
            try:
                if safe_enqueue_generation_task(
                    task,
                    generation_task_repository,
                    user_id=user_id,
                    log_prefix="[生成任务]",
                    log_task_status=True,
                ):
                    created_tasks.append(task)
                else:
                    failed_tasks.append(task)
            except UserPendingLimitExceeded as _e:
                # 兜底：如果预检查后又并发提交了，在这里也拦住
                failed_tasks.append(task)
                if not created_tasks:
                    raise HTTPException(
                        status_code=429,
                        detail="您的待处理任务过多，请等待完成后再提交",
                    ) from _e
                break
            except GlobalQueueFull as _e:
                failed_tasks.append(task)
                if not created_tasks:
                    raise HTTPException(
                        status_code=503,
                        detail="系统繁忙，请稍后再试",
                    ) from _e
                break
    except HTTPException:
        raise
    except Exception as e:
        logger.error("[生成任务] 创建失败: %s", e, exc_info=True)
        raise HTTPException(status_code=500, detail="创建生成任务失败，请稍后重试或查看任务日志") from e

    items = [_to_generation_task_response(t) for t in created_tasks + failed_tasks]
    return BatchGenerationTaskResponse(items=items, total=len(items))


@router.get("/tasks", response_model=ListGenerationTasksResponse)
def list_generation_tasks(
    authenticated_user: AuthenticatedUser = Depends(get_current_user),
    generation_task_repository: Any = Depends(get_generation_task_repository),
) -> ListGenerationTasksResponse:
    """用户级生成任务列表（跨 project）。"""
    tasks = generation_task_repository.list_by_user(authenticated_user.user.id)
    items = [_to_generation_task_response(task) for task in tasks]
    return ListGenerationTasksResponse(items=items)


@router.get("/tasks/{task_id}", response_model=GenerationTaskResponse)
def get_generation_task(
    task_id: str,
    authenticated_user: AuthenticatedUser = Depends(get_current_user),
    generation_task_repository: Any = Depends(get_generation_task_repository),
    project_repository: Any = Depends(get_project_repository),
) -> GenerationTaskResponse:
    use_case = GetGenerationTaskUseCase(generation_task_repository)
    task = use_case.execute(task_id)
    if task is None:
        raise HTTPException(status_code=404, detail=f"GenerationTask {task_id} not found")
    if task.project_id:
        check_project_access(task.project_id, authenticated_user.user.id, project_repository)
    return _to_generation_task_response(task)


@router.get("/tasks/{task_id}/results", response_model=ListGeneratedVideosResponse)
def list_generation_results(
    task_id: str,
    authenticated_user: AuthenticatedUser = Depends(get_current_user),
    generation_task_repository: Any = Depends(get_generation_task_repository),
    generated_video_repository: Any = Depends(get_generated_video_repository),
    project_repository: Any = Depends(get_project_repository),
    storage_service: OSSStorageService = Depends(get_storage_service),
) -> ListGeneratedVideosResponse:
    task = generation_task_repository.get(task_id)
    if task is None:
        raise HTTPException(status_code=404, detail=f"GenerationTask {task_id} not found")
    if task.project_id:
        check_project_access(task.project_id, authenticated_user.user.id, project_repository)
    use_case = ListGeneratedVideosByTaskUseCase(generated_video_repository)
    items = use_case.execute(task_id)
    responses = []
    for item in items:
        download_url = storage_service.get_download_url(item.file_url, expires_seconds=86400)
        responses.append(_to_generated_video_response(item, download_url=download_url))
    return ListGeneratedVideosResponse(items=responses)


@router.post("/tasks/{task_id}/retry", response_model=GenerationTaskResponse)
def retry_generation_task(
    task_id: str,
    authenticated_user: AuthenticatedUser = Depends(get_current_user),
    generation_task_repository: Any = Depends(get_generation_task_repository),
) -> GenerationTaskResponse:
    """简化重试：通过 task_id 直接重试失败任务。"""
    task = generation_task_repository.get(task_id)
    if task is None:
        raise HTTPException(status_code=404, detail="Generation task not found")
    if task.created_by_user_id and task.created_by_user_id != authenticated_user.user.id:
        raise HTTPException(status_code=403, detail="Access denied to this task")
    status_val = task.status.value if hasattr(task.status, "value") else str(task.status)
    if status_val != "failed":
        raise HTTPException(status_code=409, detail="Only failed tasks can be retried")

    user_id = authenticated_user.user.id
    # 预检查：创建前判断，>= 上限就拒绝
    user_pending = generation_task_repository.count_pending_by_user(user_id)
    global_pending = generation_task_repository.count_pending_total()
    if user_pending >= USER_PENDING_LIMIT:
        raise HTTPException(
            status_code=429,
            detail=f"您的待处理任务过多（当前 {user_pending}/{USER_PENDING_LIMIT}），请等待完成后再提交",
        )
    if global_pending >= GLOBAL_PENDING_LIMIT:
        raise HTTPException(
            status_code=503,
            detail="系统繁忙，请稍后再试",
        )

    use_case = CreateGenerationTaskUseCase(generation_task_repository)
    retried = use_case.execute(
        CreateGenerationTaskCommand(
            project_id=task.project_id,
            asset_library_id=task.asset_library_id,
            strategy_id=task.strategy_id,
            voice_library_id=task.voice_library_id,
            template_id=task.template_id,
            asset_ids=task.asset_ids,
            title_ids=task.title_ids,
            voice_ids=task.voice_ids,
            created_by_user_id=user_id,
            source_edit_plan_id=task.source_edit_plan_id or "",
            asset_select_mode=getattr(task, "asset_select_mode", ""),
            video_title=getattr(task, "video_title", ""),
            resolution=getattr(task, "resolution", ""),
        )
    )
    try:
        if not safe_enqueue_generation_task(
            retried,
            generation_task_repository,
            user_id=user_id,
            log_prefix="[生成任务]",
            log_task_status=True,
        ):
            logger.warning("[生成任务] 重试入队失败: task_id=%s", retried.id)
    except UserPendingLimitExceeded:
        raise HTTPException(
            status_code=429,
            detail="您的待处理任务过多，请等待完成后再提交",
        ) from None
    except GlobalQueueFull:
        raise HTTPException(
            status_code=503,
            detail="系统繁忙，请稍后再试",
        ) from None
    return _to_generation_task_response(retried)


@router.post("/tasks/{task_id}/cancel", response_model=GenerationTaskResponse)
def cancel_generation_task(
    task_id: str,
    authenticated_user: AuthenticatedUser = Depends(get_current_user),
    generation_task_repository: Any = Depends(get_generation_task_repository),
) -> GenerationTaskResponse:
    """取消生成任务。

    仅 pending / running 状态的任务可取消；取消后状态变为 cancelled。
    对于已在运行的 Celery 任务，标记为 cancelled 后，worker 在下次检查点会中止执行。
    """
    task = generation_task_repository.get(task_id)
    if task is None:
        raise HTTPException(status_code=404, detail="Generation task not found")

    # 权限校验
    if task.created_by_user_id and task.created_by_user_id != authenticated_user.user.id:
        raise HTTPException(status_code=403, detail="Access denied to this task")

    status_val = task.status.value if hasattr(task.status, "value") else str(task.status)

    # 终态不可取消
    if status_val in ("completed", "failed", "cancelled"):
        raise HTTPException(
            status_code=409,
            detail=f"Cannot cancel task in {status_val} status",
        )

    # 执行取消
    try:
        task.mark_cancelled()
        task.append_log(
            stage="cancelled",
            message="用户主动取消任务",
            level="INFO",
            cancelled_by=authenticated_user.user.id,
        )
        generation_task_repository.update(task)
        logger.info(
            "生成任务已取消: task_id=%s user_id=%s previous_status=%s",
            task_id,
            authenticated_user.user.id,
            status_val,
        )
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e)) from e

    return _to_generation_task_response(task)