video-gen/video-gen-api/app/services/module_generation_flow_base_service.py

from __future__ import annotations

from collections.abc import Awaitable, Callable
from datetime import datetime
from typing import Any

from fastapi import HTTPException
from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.enums.common import ModuleEventTypeEnum, ModuleStepStatusEnum
from app.models.chat_generation_task import ChatGenerationTask
from app.models.module_generation_project import ModuleGenerationProject
from app.models.module_generation_step import ModuleGenerationStep
from app.models.user import User
from app.services.generation_refund_service import mark_chat_generation_task_failed_and_refund_once
from app.enums.module_generation_flow import ModuleGenerationFlowConfig
from app.services.module_generation_step_common_service import build_step_input, build_step_output, utc_now
from app.services.resource_accounting_service import soft_delete_chat_task_resources
from app.utils.id_gen import generate_id

LogModuleEventCallable = Callable[..., Awaitable[None]]


async def get_project_for_user(
    db: AsyncSession,
    *,
    project_id: str,
    user: User,
    config: ModuleGenerationFlowConfig,
    for_update: bool = False,
    populate_existing: bool = False,
) -> ModuleGenerationProject:
    query = select(ModuleGenerationProject).where(
        ModuleGenerationProject.id == project_id,
        ModuleGenerationProject.module == config.module,
        ModuleGenerationProject.deleted_at.is_(None),
    )
    if not user.is_admin:
        query = query.where(ModuleGenerationProject.user_id == user.id)
    if populate_existing:
        query = query.execution_options(populate_existing=True)
    if for_update:
        query = query.with_for_update()
    result = await db.execute(query.limit(1))
    project = result.scalar_one_or_none()
    if not project:
        raise HTTPException(status_code=404, detail=config.project_not_found_message)
    return project


async def get_step_for_user(
    db: AsyncSession,
    *,
    project_id: str,
    step_id: str,
    user: User,
    config: ModuleGenerationFlowConfig,
    for_update: bool = False,
) -> ModuleGenerationStep:
    await get_project_for_user(db, project_id=project_id, user=user, config=config, for_update=for_update)
    query = select(ModuleGenerationStep).where(
        ModuleGenerationStep.id == step_id,
        ModuleGenerationStep.project_id == project_id,
        ModuleGenerationStep.module == config.module,
        ModuleGenerationStep.deleted_at.is_(None),
        ModuleGenerationStep.is_current == True,
    )
    if not user.is_admin:
        query = query.where(ModuleGenerationStep.user_id == user.id)
    if for_update:
        query = query.with_for_update()
    result = await db.execute(query.limit(1))
    step = result.scalar_one_or_none()
    if not step:
        raise HTTPException(status_code=404, detail=config.step_not_found_message)
    return step


async def get_current_steps(db: AsyncSession, *, project_id: str, config: ModuleGenerationFlowConfig) -> list[ModuleGenerationStep]:
    result = await db.execute(
        select(ModuleGenerationStep)
        .where(
            ModuleGenerationStep.project_id == project_id,
            ModuleGenerationStep.module == config.module,
            ModuleGenerationStep.deleted_at.is_(None),
            ModuleGenerationStep.is_current == True,
        )
        .order_by(ModuleGenerationStep.step_index.asc(), ModuleGenerationStep.created_at.asc())
    )
    return list(result.scalars().all())


async def get_current_step_by_code(
    db: AsyncSession,
    *,
    project_id: str,
    step_code: str,
    config: ModuleGenerationFlowConfig,
) -> ModuleGenerationStep | None:
    result = await db.execute(
        select(ModuleGenerationStep)
        .where(
            ModuleGenerationStep.project_id == project_id,
            ModuleGenerationStep.module == config.module,
            ModuleGenerationStep.step_code == step_code,
            ModuleGenerationStep.is_current == True,
            ModuleGenerationStep.deleted_at.is_(None),
        )
        .order_by(ModuleGenerationStep.version.desc(), ModuleGenerationStep.created_at.desc())
        .limit(1)
    )
    return result.scalar_one_or_none()


async def next_version(
    db: AsyncSession,
    *,
    project_id: str,
    step_code: str,
    config: ModuleGenerationFlowConfig,
) -> int:
    result = await db.execute(
        select(func.max(ModuleGenerationStep.version)).where(
            ModuleGenerationStep.project_id == project_id,
            ModuleGenerationStep.module == config.module,
            ModuleGenerationStep.step_code == step_code,
        )
    )
    return int(result.scalar_one_or_none() or 0) + 1


async def create_module_step(
    db: AsyncSession,
    *,
    project: ModuleGenerationProject,
    step_code: str,
    config: ModuleGenerationFlowConfig,
    log_module_event: LogModuleEventCallable,
    status: str = ModuleStepStatusEnum.PENDING.value,
    parent_step_id: str | None = None,
    source_step_id: str | None = None,
    chat_task_id: str | None = None,
    input_data: dict[str, Any] | None = None,
    output_data: dict[str, Any] | None = None,
) -> ModuleGenerationStep:
    version = await next_version(db, project_id=project.id, step_code=step_code, config=config)
    step = ModuleGenerationStep(
        id=generate_id(),
        project_id=project.id,
        user_id=project.user_id,
        module=project.module,
        step_index=config.step_index_map[step_code],
        step_code=step_code,
        status=status,
        version=version,
        is_current=True,
        parent_step_id=parent_step_id,
        source_step_id=source_step_id,
        chat_task_id=chat_task_id,
        input_json=(
            build_step_input(
                step_code=step_code,
                payload=input_data,
                source_step_id=source_step_id,
                parent_step_id=parent_step_id,
                schema_version=config.step_io_schema_version,
            )
            if input_data is not None
            else None
        ),
        output_json=(
            build_step_output(
                step_code=step_code,
                status=status,
                result=output_data,
                schema_version=config.step_io_schema_version,
            )
            if output_data is not None
            else None
        ),
        started_at=utc_now() if status == ModuleStepStatusEnum.PROCESSING.value else None,
        completed_at=utc_now() if status == ModuleStepStatusEnum.COMPLETED.value else None,
    )
    db.add(step)
    project.current_step_code = step_code
    await db.flush()
    await log_module_event(
        db,
        project=project,
        step=step,
        event_type=ModuleEventTypeEnum.STEP_CREATED.value,
        detail={"step_code": step_code, "version": version},
    )
    return step


def _clear_project_final_resources_by_deleted_steps(
    project: ModuleGenerationProject,
    *,
    deleted_step_codes: set[str],
    config: ModuleGenerationFlowConfig,
) -> list[str]:
    """根据被软删的步骤，清空项目表中对应的最终资源字段。

    第 3 步图片生成被软删时，下游第 5 步视频也基于旧图片失效，所以图片、视频、封面都要清空。
    第 5 步视频生成被软删时，只清空视频和封面，保留第 3 步图片结果。
    """
    cleared_fields: list[str] = []

    def clear_field(field_name: str) -> None:
        if getattr(project, field_name, None):
            cleared_fields.append(field_name)
        setattr(project, field_name, None)

    if config.image_generate_step_code in deleted_step_codes:
        clear_field("final_image_url")
        clear_field("final_video_url")
        clear_field("final_video_cover_url")
    elif config.video_generate_step_code in deleted_step_codes:
        clear_field("final_video_url")
        clear_field("final_video_cover_url")

    return cleared_fields


async def soft_delete_steps_from_index(
    db: AsyncSession,
    *,
    project: ModuleGenerationProject,
    start_index: int,
    config: ModuleGenerationFlowConfig,
    log_module_event: LogModuleEventCallable,
    deleted_at: datetime | None = None,
    refund_unfinished: bool = True,
    release_stats: dict[str, int] | None = None,
) -> list[ModuleGenerationStep]:
    deleted_at = deleted_at or utc_now()
    result = await db.execute(
        select(ModuleGenerationStep)
        .where(
            ModuleGenerationStep.project_id == project.id,
            ModuleGenerationStep.module == config.module,
            ModuleGenerationStep.is_current == True,
            ModuleGenerationStep.deleted_at.is_(None),
            ModuleGenerationStep.step_index >= start_index,
        )
        .with_for_update()
    )
    steps = list(result.scalars().all())
    deleted_step_codes = {step.step_code for step in steps}
    cleared_project_fields = _clear_project_final_resources_by_deleted_steps(
        project,
        deleted_step_codes=deleted_step_codes,
        config=config,
    )

    for step in steps:
        step.is_current = False
        step.deleted_at = deleted_at
        if step.chat_task_id:
            chat_result = await db.execute(
                select(ChatGenerationTask)
                .where(ChatGenerationTask.id == step.chat_task_id, ChatGenerationTask.deleted_at.is_(None))
                .with_for_update()
                .limit(1)
            )
            chat_task = chat_result.scalar_one_or_none()
            if chat_task:
                if chat_task.status == "completed":
                    released_size = await soft_delete_chat_task_resources(db, chat_task.id, deleted_at=deleted_at)
                    if release_stats is not None:
                        release_stats["released_size_bytes"] = int(release_stats.get("released_size_bytes", 0)) + int(released_size or 0)
                elif refund_unfinished and chat_task.status != "failed":
                    await mark_chat_generation_task_failed_and_refund_once(
                        db,
                        task=chat_task,
                        error_message=config.cancel_chat_task_error_message,
                        pipeline_stage="failed",
                    )
                chat_task.deleted_at = deleted_at
    if steps:
        await log_module_event(
            db,
            project=project,
            event_type=ModuleEventTypeEnum.SOFT_DELETE_STEPS.value,
            message=f"软删除第 {start_index} 步及之后的旧子任务",
            detail={
                "step_ids": [step.id for step in steps],
                "step_codes": [step.step_code for step in steps],
                "cleared_project_fields": cleared_project_fields,
            },
        )
    return steps


async def chat_tasks_by_id(db: AsyncSession, steps: list[ModuleGenerationStep]) -> dict[str, ChatGenerationTask]:
    ids = [step.chat_task_id for step in steps if step.chat_task_id]
    if not ids:
        return {}
    result = await db.execute(select(ChatGenerationTask).where(ChatGenerationTask.id.in_(ids)))
    return {task.id: task for task in result.scalars().all()}