xiaoxia-saas/packages/adapters/sqlalchemy_impl/voice_clone_profile_repository.py

"""SQLAlchemy implementation of VoiceCloneProfileRepository."""

from __future__ import annotations

from typing import Optional

from sqlalchemy.orm import Session

from packages.adapters.sqlalchemy_impl.models import VoiceCloneProfileModel
from packages.domain.voice_clone_profile import VoiceCloneProfile, VoiceCloneStatus


class SQLAlchemyVoiceCloneProfileRepository:
    """SQLAlchemy 音色克隆档案仓储。"""

    def __init__(self, session: Session) -> None:
        self.session = session

    def create(self, profile: VoiceCloneProfile) -> VoiceCloneProfile:
        model = VoiceCloneProfileModel(
            id=profile.id,
            user_id=profile.user_id,
            name=profile.name,
            description=profile.description,
            source_audio_url=profile.source_audio_url,
            voice_id=profile.voice_id,
            voice_model=profile.voice_model,
            language=profile.language,
            gender=profile.gender,
            status=profile.status,
            error_message=profile.error_message,
            retry_count=profile.retry_count,
            max_retries=profile.max_retries,
            metadata_=profile.metadata,
        )
        self.session.add(model)
        self.session.commit()
        self.session.refresh(model)
        return self._model_to_entity(model)

    def get(self, profile_id: str) -> Optional[VoiceCloneProfile]:
        model = (
            self.session.query(VoiceCloneProfileModel)
            .filter(
                VoiceCloneProfileModel.id == profile_id,
                VoiceCloneProfileModel.status != "deleted",
            )
            .first()
        )
        if model is None:
            return None
        return self._model_to_entity(model)

    def update(self, profile: VoiceCloneProfile) -> VoiceCloneProfile:
        model = self.session.query(VoiceCloneProfileModel).filter(VoiceCloneProfileModel.id == profile.id).first()
        if model is None:
            raise ValueError(f"VoiceCloneProfile {profile.id} not found")
        model.name = profile.name
        model.description = profile.description
        model.source_audio_url = profile.source_audio_url
        model.voice_id = profile.voice_id
        model.voice_model = profile.voice_model
        model.language = profile.language
        model.gender = profile.gender
        model.status = profile.status
        model.error_message = profile.error_message
        model.retry_count = profile.retry_count
        model.max_retries = profile.max_retries
        model.metadata_ = profile.metadata
        self.session.commit()
        self.session.refresh(model)
        return self._model_to_entity(model)

    def delete(self, profile_id: str) -> bool:
        model = self.session.query(VoiceCloneProfileModel).filter(VoiceCloneProfileModel.id == profile_id).first()
        if model is None:
            return False
        model.status = "deleted"
        self.session.commit()
        return True

    def list_by_user(
        self,
        user_id: str,
        *,
        status: Optional[str] = None,
        limit: int = 50,
        offset: int = 0,
    ) -> list[VoiceCloneProfile]:
        query = self.session.query(VoiceCloneProfileModel).filter(
            VoiceCloneProfileModel.user_id == user_id,
            VoiceCloneProfileModel.status != "deleted",
        )
        if status:
            query = query.filter(VoiceCloneProfileModel.status == status)
        query = query.order_by(VoiceCloneProfileModel.created_at.desc())
        models = query.offset(offset).limit(limit).all()
        return [self._model_to_entity(m) for m in models]

    def count_by_user(self, user_id: str, *, status: Optional[str] = None) -> int:
        query = self.session.query(VoiceCloneProfileModel).filter(
            VoiceCloneProfileModel.user_id == user_id,
            VoiceCloneProfileModel.status != "deleted",
        )
        if status:
            query = query.filter(VoiceCloneProfileModel.status == status)
        return query.count()

    def find_by_voice_id(self, voice_id: str) -> Optional[VoiceCloneProfile]:
        model = (
            self.session.query(VoiceCloneProfileModel)
            .filter(
                VoiceCloneProfileModel.voice_id == voice_id,
                VoiceCloneProfileModel.status != "deleted",
            )
            .first()
        )
        if model is None:
            return None
        return self._model_to_entity(model)

    def find_profile_ids_by_voice_ids(self, voice_ids: list[str]) -> dict[str, str]:
        """批量查询 voice_id → profile_id 映射。用于填充统一列表的 voice_clone_profile_id。"""
        if not voice_ids:
            return {}
        rows = (
            self.session.query(
                VoiceCloneProfileModel.voice_id,
                VoiceCloneProfileModel.id,
            )
            .filter(
                VoiceCloneProfileModel.voice_id.in_(voice_ids),
                VoiceCloneProfileModel.status != "deleted",
            )
            .all()
        )
        return {voice_id: profile_id for voice_id, profile_id in rows}

    def cleanup_stale_processing(self, timeout_minutes: int = 10) -> int:
        """清理超时卡在 processing 的克隆档案。

        worker 重启、Celery 任务丢失或 OOM 被杀时，processing 档案会永久卡住。
        updated_at < NOW() - timeout_minutes 的 processing 记录，标记为 failed
        并附带明确错误信息，用户可在前端点击「重试」。

        Args:
            timeout_minutes: 超时分钟数，默认 10 分钟（正常克隆 < 5 分钟）

        Returns:
            清理的记录数
        """
        from datetime import datetime, timedelta, UTC

        cutoff = datetime.now(UTC) - timedelta(minutes=timeout_minutes)
        models = (
            self.session.query(VoiceCloneProfileModel)
            .filter(
                VoiceCloneProfileModel.status == "processing",
                VoiceCloneProfileModel.updated_at < cutoff,
            )
            .all()
        )
        count = 0
        for model in models:
            model.status = "failed"
            model.error_message = f"克隆任务执行超时（超过 {timeout_minutes} 分钟未更新，可能因服务重启中断），请重试"
            count += 1
        if count > 0:
            self.session.commit()
        return count

    @staticmethod
    def _model_to_entity(model: VoiceCloneProfileModel) -> VoiceCloneProfile:
        return VoiceCloneProfile(
            id=model.id,
            user_id=model.user_id,
            name=model.name,
            description=model.description or "",
            source_audio_url=model.source_audio_url or "",
            voice_id=model.voice_id or "",
            voice_model=model.voice_model or "",
            language=model.language or "zh-CN",
            gender=model.gender or "unknown",
            status=VoiceCloneStatus(model.status) if model.status else VoiceCloneStatus.PENDING,
            error_message=model.error_message or "",
            retry_count=model.retry_count or 0,
            max_retries=model.max_retries or 3,
            metadata=model.metadata_ or {},
            created_at=model.created_at,
            updated_at=model.updated_at,
        )