xiaoxia-saas/packages/domain/generation_task.py

"""GenerationTask 领域模型 — 视频生成任务.

状态机：
  pending → running → completed
              ↘ failed → pending (重试)
       ↘ cancelled
"""

from __future__ import annotations

import sys
from dataclasses import dataclass, field
from datetime import datetime, timezone

if sys.version_info >= (3, 11):
    from enum import StrEnum
else:
    from enum import Enum

    class StrEnum(str, Enum):
        pass


from uuid import uuid4


class GenerationTaskStatus(StrEnum):
    """生成任务状态枚举。"""

    PENDING = "pending"
    """待处理（任务已创建，等待执行）"""

    RUNNING = "running"
    """运行中（正在生成视频）"""

    COMPLETED = "completed"
    """已完成（视频生成成功）"""

    FAILED = "failed"
    """失败（生成失败）"""

    CANCELLED = "cancelled"
    """已取消（用户取消或系统取消）"""


# 终态集合
TERMINAL_STATUSES = frozenset(
    {GenerationTaskStatus.COMPLETED, GenerationTaskStatus.FAILED, GenerationTaskStatus.CANCELLED}
)

# 合法状态转换
_VALID_TRANSITIONS: dict[GenerationTaskStatus, set[GenerationTaskStatus]] = {
    GenerationTaskStatus.PENDING: {
        GenerationTaskStatus.RUNNING,
        GenerationTaskStatus.FAILED,
        GenerationTaskStatus.CANCELLED,
    },
    GenerationTaskStatus.RUNNING: {
        GenerationTaskStatus.COMPLETED,
        GenerationTaskStatus.FAILED,
        GenerationTaskStatus.CANCELLED,
    },
    GenerationTaskStatus.FAILED: {GenerationTaskStatus.PENDING},  # 重试回到 pending
}


@dataclass(slots=True)
class GenerationTask:
    id: str
    project_id: str
    asset_library_id: str
    strategy_id: str = ""
    voice_library_id: str = ""
    template_id: str = ""
    asset_ids: list[str] = field(default_factory=list)
    title_ids: list[str] = field(default_factory=list)
    voice_ids: list[str] = field(default_factory=list)
    status: GenerationTaskStatus = GenerationTaskStatus.PENDING
    progress: float = 0.0
    result_count: int = 0
    error_message: str = ""
    started_at: datetime | None = None
    completed_at: datetime | None = None
    source_edit_plan_id: str = ""
    created_by_user_id: str = ""
    asset_select_mode: str = ""
    batch_id: str = ""
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))

    @classmethod
    def create(
        cls,
        project_id: str,
        asset_library_id: str,
        *,
        strategy_id: str = "",
        voice_library_id: str = "",
        template_id: str = "",
        asset_ids: list[str] | None = None,
        title_ids: list[str] | None = None,
        voice_ids: list[str] | None = None,
        created_by_user_id: str = "",
        source_edit_plan_id: str = "",
        asset_select_mode: str = "",
        batch_id: str = "",
    ) -> "GenerationTask":
        if not project_id.strip() and not template_id.strip():
            raise ValueError("project_id 或 template_id 至少需要提供一个")
        if not asset_library_id.strip() and not (asset_ids or title_ids or voice_ids):
            raise ValueError("asset_library_id 或 asset_ids/title_ids/voice_ids 至少需要提供一个")
        return cls(
            id=uuid4().hex,
            project_id=project_id.strip(),
            asset_library_id=asset_library_id.strip(),
            strategy_id=strategy_id.strip(),
            voice_library_id=voice_library_id.strip(),
            template_id=template_id.strip(),
            asset_ids=list(asset_ids) if asset_ids else [],
            title_ids=list(title_ids) if title_ids else [],
            voice_ids=list(voice_ids) if voice_ids else [],
            created_by_user_id=created_by_user_id.strip(),
            source_edit_plan_id=source_edit_plan_id.strip(),
            asset_select_mode=asset_select_mode,
            batch_id=batch_id,
        )

    # ── 状态查询 ────────────────────────────────────────────────────────────

    @property
    def is_terminal(self) -> bool:
        """是否处于终态（completed / failed / cancelled）。"""
        return self.status in TERMINAL_STATUSES

    @property
    def is_completed(self) -> bool:
        """是否已完成。"""
        return self.status == GenerationTaskStatus.COMPLETED

    @property
    def is_failed(self) -> bool:
        """是否失败。"""
        return self.status == GenerationTaskStatus.FAILED

    @property
    def is_running(self) -> bool:
        """是否运行中。"""
        return self.status == GenerationTaskStatus.RUNNING

    # ── 状态转换 ────────────────────────────────────────────────────────────

    def transition_to(self, new_status: GenerationTaskStatus | str) -> None:
        """执行状态转换。

        Args:
            new_status: 目标状态

        Raises:
            ValueError: 非法状态转换
        """
        if isinstance(new_status, str):
            try:
                new_status = GenerationTaskStatus(new_status)
            except ValueError:
                raise ValueError(f"无效状态: {new_status}")

        allowed = _VALID_TRANSITIONS.get(self.status, set())
        if new_status not in allowed:
            raise ValueError(
                f"非法状态转换: {self.status.value} → {new_status.value}，"
                f"允许: {{{', '.join(sorted(s.value for s in allowed))}}}"
            )

        self.status = new_status

    def mark_processing(self) -> None:
        """标记为处理中（pending → running）。

        设置 started_at，清除 error_message。

        Raises:
            ValueError: 当前状态不允许转换到 running
        """
        self.transition_to(GenerationTaskStatus.RUNNING)
        self.started_at = datetime.now(timezone.utc)
        self.error_message = ""

    def mark_completed(self, result_count: int = 1) -> None:
        """标记为已完成（running → completed）。

        设置 completed_at、progress=100.0、result_count，清除 error_message。

        Args:
            result_count: 生成的视频数量，默认为 1

        Raises:
            ValueError: 当前状态不允许转换到 completed
        """
        self.transition_to(GenerationTaskStatus.COMPLETED)
        self.completed_at = datetime.now(timezone.utc)
        self.progress = 100.0
        self.result_count = result_count
        self.error_message = ""

    def mark_failed(self, error_message: str) -> None:
        """标记为失败（pending / running → failed）。

        设置 error_message、completed_at。

        Args:
            error_message: 错误信息

        Raises:
            ValueError: 当前状态不允许转换到 failed
        """
        self.transition_to(GenerationTaskStatus.FAILED)
        self.error_message = error_message
        self.completed_at = datetime.now(timezone.utc)

    def mark_cancelled(self) -> None:
        """标记为已取消（pending / running → cancelled）。

        设置 completed_at。

        Raises:
            ValueError: 当前状态不允许转换到 cancelled
        """
        self.transition_to(GenerationTaskStatus.CANCELLED)
        self.completed_at = datetime.now(timezone.utc)

    def mark_pending_from_failed(self) -> None:
        """从失败状态重置为待处理（用于重试）。

        清除 error_message、started_at、completed_at、progress。

        Raises:
            ValueError: 当前状态不是 failed
        """
        if self.status != GenerationTaskStatus.FAILED:
            raise ValueError(
                f"只有 failed 状态的任务可以重置为 pending，当前状态: {self.status.value}"
            )
        self.transition_to(GenerationTaskStatus.PENDING)
        self.error_message = ""
        self.started_at = None
        self.completed_at = None
        self.progress = 0.0
        self.result_count = 0