xiaoxia-saas/packages/domain/generation_task.py

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

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

from __future__ import annotations

import json
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 = ""
    error_info: dict = field(default_factory=dict)
    retry_count: int = 0
    auto_retry_enabled: bool = False
    auto_retry_max: int = 0
    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 = ""
    video_title: str = ""
    resolution: str = ""
    bgm_config: dict = field(default_factory=dict)
    logs: str = "[]"
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    updated_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 = "",
        video_title: str = "",
        resolution: str = "",
        bgm_config: dict | None = None,
        auto_retry_enabled: bool = False,
        auto_retry_max: int = 0,
    ) -> "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,
            video_title=video_title.strip(),
            resolution=resolution.strip(),
            bgm_config=dict(bgm_config) if bgm_config else {},
            auto_retry_enabled=auto_retry_enabled,
            auto_retry_max=auto_retry_max,
        )

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

    @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 as _e:
                raise ValueError(f"无效状态: {new_status}") from _e

        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, error_info: dict | None = None) -> None:
        """标记为失败（pending / running → failed）。

        设置 error_message、error_info、completed_at。

        Args:
            error_message: 错误信息
            error_info: 结构化错误信息（error_type, stack_trace, stage, failed_at等）

        Raises:
            ValueError: 当前状态不允许转换到 failed
        """
        self.transition_to(GenerationTaskStatus.FAILED)
        self.error_message = error_message
        self.completed_at = datetime.now(timezone.utc)
        if error_info is not None:
            self.error_info = error_info
        else:
            self.error_info = {
                "error_type": "UnknownError",
                "message": error_message,
                "failed_at": datetime.now(timezone.utc).isoformat(),
            }

    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)

    # ── 日志辅助 ────────────────────────────────────────────────────────────

    _MAX_LOGS = 200

    def append_log(self, stage: str, message: str, level: str = "INFO", **kwargs) -> None:
        """追加一条结构化日志到 logs 字段。

        Args:
            stage: 阶段名称（如 "接收任务"、"下载素材"、"渲染"）
            message: 日志消息
            level: 日志级别（INFO / WARN / ERROR）
            **kwargs: 额外字段（如 asset_id、duration 等）
        """
        try:
            entries = json.loads(self.logs) if self.logs else []
        except (json.JSONDecodeError, TypeError):
            entries = []
        entry = {
            "ts": datetime.now(timezone.utc).isoformat(),
            "level": level,
            "stage": stage,
            "message": message,
            **kwargs,
        }
        entries.append(entry)
        # 限制最多保留 _MAX_LOGS 条，防止字段过大
        if len(entries) > self._MAX_LOGS:
            entries = entries[-self._MAX_LOGS :]
        self.logs = json.dumps(entries, ensure_ascii=False)

    def get_logs(self) -> list[dict]:
        """解析 logs 字段为 list[dict]。"""
        try:
            return json.loads(self.logs) if self.logs else []
        except (json.JSONDecodeError, TypeError):
            return []

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

        清除 error_message、error_info、started_at、completed_at、progress，
        递增 retry_count。

        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.error_info = {}
        self.started_at = None
        self.completed_at = None
        self.progress = 0.0
        self.result_count = 0
        self.retry_count += 1