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

import json
from datetime import UTC, datetime

from sqlalchemy.orm import Session

from packages.adapters.sqlalchemy_impl.models import AssetLibraryModel, AssetModel, AssetTagModel
from packages.domain import Asset, AssetStatus, ClassificationStatus


class SQLAlchemyAssetRepository:
    def __init__(self, session: Session):
        self.session = session

    def find_by_library(
        self,
        library_id: str,
        skip: int = 0,
        limit: int = 100,
        status: list[str] | None = None,
    ) -> list[Asset]:
        query = self.session.query(AssetModel).filter(AssetModel.asset_library_id == library_id)
        if status:
            query = query.filter(AssetModel.status.in_(status))
        models = query.order_by(AssetModel.created_at.desc()).offset(skip).limit(limit).all()
        return [self._to_domain(model) for model in models]

    def find_by_project(
        self,
        project_id: str,
        skip: int = 0,
        limit: int = 100,
        status: list[str] | None = None,
    ) -> list[Asset]:
        query = self.session.query(AssetModel).filter(AssetModel.project_id == project_id)
        if status:
            query = query.filter(AssetModel.status.in_(status))
        models = query.order_by(AssetModel.created_at.desc()).offset(skip).limit(limit).all()
        return [self._to_domain(model) for model in models]

    def find_by_library_and_file_type(
        self,
        library_id: str,
        file_type: str,
        skip: int = 0,
        limit: int = 100,
        status: list[str] | None = None,
    ) -> list[Asset]:
        query = self.session.query(AssetModel).filter(
            AssetModel.asset_library_id == library_id, AssetModel.file_type == file_type
        )
        if status:
            query = query.filter(AssetModel.status.in_(status))
        models = query.order_by(AssetModel.created_at.desc()).offset(skip).limit(limit).all()
        return [self._to_domain(model) for model in models]

    def count_by_library_and_file_type(
        self,
        library_id: str,
        file_type: str,
        status: list[str] | None = None,
    ) -> int:
        query = self.session.query(AssetModel).filter(
            AssetModel.asset_library_id == library_id, AssetModel.file_type == file_type
        )
        if status:
            query = query.filter(AssetModel.status.in_(status))
        return query.count()

    def find_by_project_and_file_type(
        self,
        project_id: str,
        file_type: str,
        skip: int = 0,
        limit: int = 100,
        status: list[str] | None = None,
    ) -> list[Asset]:
        query = self.session.query(AssetModel).filter(
            AssetModel.project_id == project_id, AssetModel.file_type == file_type
        )
        if status:
            query = query.filter(AssetModel.status.in_(status))
        models = query.order_by(AssetModel.created_at.desc()).offset(skip).limit(limit).all()
        return [self._to_domain(model) for model in models]

    def count_by_project_and_file_type(
        self,
        project_id: str,
        file_type: str,
        status: list[str] | None = None,
    ) -> int:
        query = self.session.query(AssetModel).filter(
            AssetModel.project_id == project_id, AssetModel.file_type == file_type
        )
        if status:
            query = query.filter(AssetModel.status.in_(status))
        return query.count()

    def find_by_id(self, asset_id: str) -> Asset | None:
        model = self.session.query(AssetModel).filter(AssetModel.id == asset_id).first()
        if model is None:
            return None
        return self._to_domain(model)

    def find_by_ids(self, asset_ids: list[str]) -> list[Asset]:
        """批量查询素材（单次 SQL IN 查询，避免 N+1）。"""
        if not asset_ids:
            return []
        models = self.session.query(AssetModel).filter(AssetModel.id.in_(asset_ids)).all()
        return [self._to_domain(m) for m in models]

    def get(self, asset_id: str) -> Asset | None:
        return self.find_by_id(asset_id)

    def create(self, asset: Asset) -> Asset:
        now = datetime.now(UTC)
        model = AssetModel(
            id=asset.id,
            project_id=asset.project_id,
            asset_library_id=asset.library_id,
            name=asset.name,
            file_type=(asset.mime_type.split("/")[0] if "/" in asset.mime_type else asset.mime_type),
            file_size=asset.file_size,
            file_url=asset.storage_key,
            storage_key=asset.storage_key,
            thumbnail_url=asset.thumbnail_url,
            duration=asset.duration,
            width=asset.width,
            height=asset.height,
            fps=asset.fps,
            codec=asset.codec,
            status=(asset.status.value if hasattr(asset.status, "value") else str(asset.status)),
            classification_status=(
                asset.classification_status.value
                if hasattr(asset.classification_status, "value")
                else str(asset.classification_status)
            ),
            classification_result=(json.dumps(asset.metadata) if asset.metadata else None),
            quality_score=asset.quality_score,
            uploaded_by_user_id=asset.uploaded_by_user_id or "system",
            file_hash=asset.file_hash or None,
            client_upload_id=asset.client_upload_id or None,
            created_at=asset.created_at,
            updated_at=now,
        )
        self.session.add(model)
        self.session.flush()
        self._sync_asset_tags(asset.id, asset.tag_ids)
        # Issue #1776: 自动维护素材库计数（同事务内原子更新）
        if asset.library_id and (getattr(asset.status, "value", str(asset.status)) != "deleted"):
            from sqlalchemy import func

            self.session.query(AssetLibraryModel).filter(AssetLibraryModel.id == asset.library_id).update(
                {
                    AssetLibraryModel.asset_count: func.coalesce(AssetLibraryModel.asset_count, 0) + 1,
                    AssetLibraryModel.total_size: func.coalesce(AssetLibraryModel.total_size, 0) + asset.file_size,
                }
            )
        self.session.commit()
        return asset

    def update(self, asset: Asset) -> Asset:
        model = self.session.query(AssetModel).filter(AssetModel.id == asset.id).first()
        if model is None:
            raise ValueError(f"Asset {asset.id} not found")
        model.name = asset.name
        model.file_size = asset.file_size
        model.file_url = asset.storage_key
        model.storage_key = asset.storage_key
        model.thumbnail_url = asset.thumbnail_url
        model.duration = asset.duration
        model.width = asset.width
        model.height = asset.height
        model.fps = asset.fps
        model.codec = asset.codec
        model.status = asset.status.value if hasattr(asset.status, "value") else str(asset.status)
        model.classification_status = (
            asset.classification_status.value
            if hasattr(asset.classification_status, "value")
            else str(asset.classification_status)
        )
        model.classification_result = json.dumps(asset.metadata) if asset.metadata else None
        model.quality_score = asset.quality_score
        model.uploaded_by_user_id = asset.uploaded_by_user_id or model.uploaded_by_user_id
        model.file_hash = asset.file_hash or model.file_hash
        if getattr(model, "client_upload_id", None) is None and asset.client_upload_id:
            model.client_upload_id = asset.client_upload_id
        model.updated_at = datetime.now(UTC)
        self.session.flush()
        self._sync_asset_tags(asset.id, asset.tag_ids)
        self.session.commit()
        return asset

    def delete(self, asset_id: str) -> bool:
        model = self.session.query(AssetModel).filter(AssetModel.id == asset_id).first()
        if model:
            library_id = model.asset_library_id
            file_size = model.file_size or 0
            # 只统计非 deleted 状态的素材
            was_counted = model.status != "deleted"
            self.session.delete(model)
            # Issue #1776: 自动维护素材库计数
            if library_id and was_counted:
                from sqlalchemy import case, func

                self.session.query(AssetLibraryModel).filter(AssetLibraryModel.id == library_id).update(
                    {
                        AssetLibraryModel.asset_count: case(
                            (func.coalesce(AssetLibraryModel.asset_count, 0) - 1 < 0, 0),
                            else_=func.coalesce(AssetLibraryModel.asset_count, 0) - 1,
                        ),
                        AssetLibraryModel.total_size: case(
                            (func.coalesce(AssetLibraryModel.total_size, 0) - file_size < 0, 0),
                            else_=func.coalesce(AssetLibraryModel.total_size, 0) - file_size,
                        ),
                    }
                )
            self.session.commit()
            return True
        return False

    def batch_delete(self, asset_ids: list[str]) -> int:
        """批量删除素材（软删除，标记 status=deleted），返回实际影响数量。"""
        if not asset_ids:
            return 0
        from datetime import datetime

        now = datetime.now(UTC)
        # 先查询待删除素材的库分布（用于更新计数）
        to_delete = (
            self.session.query(AssetModel.asset_library_id, AssetModel.file_size)
            .filter(AssetModel.id.in_(asset_ids), AssetModel.status != "deleted")
            .all()
        )
        if not to_delete:
            return 0
        # 按库分组统计
        library_deltas: dict[str, tuple[int, int]] = {}  # library_id -> (count_delta, size_delta)
        for lib_id, size in to_delete:
            if lib_id not in library_deltas:
                library_deltas[lib_id] = (0, 0)
            c, s = library_deltas[lib_id]
            library_deltas[lib_id] = (c + 1, s + (size or 0))
        # 执行软删除
        count = (
            self.session.query(AssetModel)
            .filter(AssetModel.id.in_(asset_ids), AssetModel.status != "deleted")
            .update({AssetModel.status: "deleted", AssetModel.updated_at: now}, synchronize_session=False)
        )
        # Issue #1776: 自动维护各素材库计数
        if library_deltas:
            from sqlalchemy import case, func

            for lib_id, (count_delta, size_delta) in library_deltas.items():
                self.session.query(AssetLibraryModel).filter(AssetLibraryModel.id == lib_id).update(
                    {
                        AssetLibraryModel.asset_count: case(
                            (func.coalesce(AssetLibraryModel.asset_count, 0) - count_delta < 0, 0),
                            else_=func.coalesce(AssetLibraryModel.asset_count, 0) - count_delta,
                        ),
                        AssetLibraryModel.total_size: case(
                            (func.coalesce(AssetLibraryModel.total_size, 0) - size_delta < 0, 0),
                            else_=func.coalesce(AssetLibraryModel.total_size, 0) - size_delta,
                        ),
                    }
                )
        self.session.commit()
        return count

    def batch_update_metadata(self, asset_ids: list[str], metadata_patch: dict[str, object]) -> int:
        """批量更新素材 metadata（合并 patch），返回实际影响数量。"""
        if not asset_ids:
            return 0
        from datetime import datetime

        now = datetime.now(UTC)
        # 逐条读取 + 合并 + 更新，保证 JSON 合并正确
        models = self.session.query(AssetModel).filter(AssetModel.id.in_(asset_ids)).all()
        count = 0
        for model in models:
            existing = {}
            if model.classification_result:
                try:
                    existing = json.loads(model.classification_result)
                except Exception:
                    existing = {}
            merged = {**existing, **metadata_patch}
            model.classification_result = json.dumps(merged, ensure_ascii=False)
            model.updated_at = now
            count += 1
        self.session.commit()
        return count

    def batch_add_tags(self, asset_ids: list[str], tag_ids: list[str]) -> int:
        """批量给素材添加标签（合并去重），返回实际影响数量。"""
        if not asset_ids or not tag_ids:
            return 0
        from datetime import datetime

        now = datetime.now(UTC)
        clean_tag_ids = list(set(tag_ids))
        count = 0
        for aid in asset_ids:
            # 查询现有标签
            existing = {
                row.tag_id
                for row in self.session.query(AssetTagModel.tag_id).filter(AssetTagModel.asset_id == aid).all()
            }
            new_tags = [t for t in clean_tag_ids if t not in existing]
            if new_tags:
                for tid in new_tags:
                    self.session.add(AssetTagModel(asset_id=aid, tag_id=tid))
                # 更新 updated_at
                self.session.query(AssetModel).filter(AssetModel.id == aid).update(
                    {AssetModel.updated_at: now}, synchronize_session=False
                )
                count += 1
        self.session.commit()
        return count

    def batch_replace_tags(self, asset_ids: list[str], tag_ids: list[str]) -> int:
        """批量替换素材标签（全量覆盖），返回实际影响数量。"""
        if not asset_ids:
            return 0
        from datetime import datetime

        now = datetime.now(UTC)
        clean_tag_ids = list(set(tag_ids))
        count = 0
        for aid in asset_ids:
            # 先删再加
            self.session.query(AssetTagModel).filter(AssetTagModel.asset_id == aid).delete(synchronize_session=False)
            for tid in clean_tag_ids:
                self.session.add(AssetTagModel(asset_id=aid, tag_id=tid))
            # 更新 updated_at
            self.session.query(AssetModel).filter(AssetModel.id == aid).update(
                {AssetModel.updated_at: now}, synchronize_session=False
            )
            count += 1
        self.session.commit()
        return count

    def count_by_project(self, project_id: str, status: list[str] | None = None) -> int:
        query = self.session.query(AssetModel).filter(AssetModel.project_id == project_id)
        if status:
            query = query.filter(AssetModel.status.in_(status))
        return query.count()

    def count_by_project_ids(self, project_ids: list[str], status: list[str] | None = None) -> int:
        if not project_ids:
            return 0
        query = self.session.query(AssetModel).filter(AssetModel.project_id.in_(project_ids))
        if status:
            query = query.filter(AssetModel.status.in_(status))
        return query.count()

    def sum_storage_by_project_ids(self, project_ids: list[str]) -> int:
        if not project_ids:
            return 0
        from sqlalchemy import func

        result = (
            self.session.query(func.coalesce(func.sum(AssetModel.file_size), 0))
            .filter(AssetModel.project_id.in_(project_ids))
            .scalar()
        )
        return int(result or 0)

    def find_ready_videos_by_user(
        self,
        user_id: str,
        *,
        limit: int = 50,
    ) -> list[Asset]:
        """查找用户上传的所有就绪视频素材。"""
        query = self.session.query(AssetModel).filter(
            AssetModel.uploaded_by_user_id == user_id,
            AssetModel.status == "ready",
            AssetModel.file_type == "video",
        )
        query = query.order_by(AssetModel.created_at.desc())
        if limit > 0:
            query = query.limit(limit)
        models = query.all()
        return [self._to_domain(m) for m in models]

    def search_candidates(
        self,
        project_id: str,
        *,
        file_type: str | None = None,
        min_quality_score: float | None = None,
        min_duration: float | None = None,
        max_duration: float | None = None,
        classification_category: str | None = None,
        tags: list[str] | None = None,
        status: str | None = None,
        limit: int = 50,
    ) -> list[Asset]:
        """按筛选条件搜索候选素材，按质量分降序排列。"""
        query = self.session.query(AssetModel).filter(
            AssetModel.project_id == project_id,
        )
        if file_type is not None:
            query = query.filter(AssetModel.file_type == file_type)
        if min_quality_score is not None:
            query = query.filter(AssetModel.quality_score >= min_quality_score)
        if min_duration is not None:
            query = query.filter(AssetModel.duration >= min_duration)
        if max_duration is not None:
            query = query.filter(AssetModel.duration <= max_duration)
        if status is not None:
            query = query.filter(AssetModel.status == status)
        if classification_category is not None:
            # classification_result 是 JSON Text，用 LIKE 匹配 category 字段
            query = query.filter(AssetModel.classification_result.like(f'%"{classification_category}"%'))
        query = query.order_by(AssetModel.quality_score.desc().nullslast())
        if limit > 0:
            query = query.limit(limit)
        models = query.all()
        candidates = [self._to_domain(m) for m in models]
        # 内存中过滤 tags（tags 存在 metadata 中）
        if tags:
            tag_set = set(tags)
            candidates = [a for a in candidates if tag_set.issubset(set(a.metadata.get("tags", [])))]
        return candidates

    def _to_domain(self, model: AssetModel) -> Asset:
        metadata = {}
        if model.classification_result:
            try:
                metadata = json.loads(model.classification_result)
            except Exception:
                metadata = {}
        mime_type = model.file_type
        if "/" not in mime_type:
            mime_type = {
                "video": "video/mp4",
                "audio": "audio/mpeg",
                "image": "image/jpeg",
            }.get(mime_type, mime_type)
        # 查询关联的 tag_ids
        tag_ids = [
            row.tag_id
            for row in self.session.query(AssetTagModel.tag_id).filter(AssetTagModel.asset_id == model.id).all()
        ]
        return Asset(
            id=model.id,
            project_id=model.project_id,
            library_id=model.asset_library_id,
            name=model.name,
            storage_key=model.storage_key or model.file_url,
            mime_type=mime_type,
            file_size=int(model.file_size or 0),
            thumbnail_url=model.thumbnail_url,
            duration=model.duration,
            width=int(model.width) if model.width is not None else None,
            height=int(model.height) if model.height is not None else None,
            fps=model.fps,
            codec=model.codec,
            status=AssetStatus(model.status),
            classification_status=ClassificationStatus(model.classification_status),
            quality_score=model.quality_score,
            uploaded_by_user_id=model.uploaded_by_user_id,
            file_hash=model.file_hash or "",
            client_upload_id=getattr(model, "client_upload_id", None) or "",
            metadata=metadata,
            tag_ids=tag_ids,
            created_at=model.created_at,
            updated_at=model.updated_at,
        )

    def _sync_asset_tags(self, asset_id: str, tag_ids: list[str]) -> None:
        """同步素材-标签关联表（全量替换）。"""
        self.session.query(AssetTagModel).filter(AssetTagModel.asset_id == asset_id).delete(synchronize_session=False)
        for tag_id in tag_ids:
            self.session.add(AssetTagModel(asset_id=asset_id, tag_id=tag_id))

    def find_by_tag_ids(
        self,
        tag_ids: list[str],
        skip: int = 0,
        limit: int = 100,
    ) -> list[Asset]:
        """查找包含所有指定标签的素材。"""
        if not tag_ids:
            return []
        from sqlalchemy import func

        # 找出同时拥有所有指定 tag_id 的 asset_id
        tag_set = set(tag_ids)
        asset_ids = (
            self.session.query(AssetTagModel.asset_id)
            .filter(AssetTagModel.tag_id.in_(tag_set))
            .group_by(AssetTagModel.asset_id)
            .having(func.count(AssetTagModel.tag_id) == len(tag_set))
            .all()
        )
        ids = [row[0] for row in asset_ids]
        if not ids:
            return []
        models = self.session.query(AssetModel).filter(AssetModel.id.in_(ids)).offset(skip).limit(limit).all()
        return [self._to_domain(m) for m in models]

    def find_by_storage_key(self, storage_key: str) -> Asset | None:
        """按 storage_key 查找素材。

        Bug #2110: 历史数据 file_url 列可能是旧路径（assets/...），新代码统一写入
        storage_key 列。双列 OR 查询，避免占位 asset 因路径错配导致 ingest 兜底新建
        第二条 READY 记录，原占位卡 PROCESSING → 前端缩略图出现后消失。
        """
        if not storage_key:
            return None
        model = (
            self.session.query(AssetModel)
            .filter((AssetModel.storage_key == storage_key) | (AssetModel.file_url == storage_key))
            .first()
        )
        if model is None:
            return None
        return self._to_domain(model)

    def find_by_library_and_file_hash(
        self,
        library_id: str,
        file_hash: str,
    ) -> Asset | None:
        """按素材库 + 文件哈希查找已有素材（去重检测）。"""
        if not file_hash:
            return None
        model = (
            self.session.query(AssetModel)
            .filter(
                AssetModel.asset_library_id == library_id,
                AssetModel.file_hash == file_hash,
            )
            .first()
        )
        if model is None:
            return None
        return self._to_domain(model)

    def find_by_library_and_client_upload_id(
        self,
        library_id: str,
        client_upload_id: str,
    ) -> Asset | None:
        """按素材库 + 客户端幂等 token 查找已有素材（complete 幂等）。"""
        if not client_upload_id:
            return None
        model = (
            self.session.query(AssetModel)
            .filter(
                AssetModel.asset_library_id == library_id,
                AssetModel.client_upload_id == client_upload_id,
            )
            .first()
        )
        if model is None:
            return None
        return self._to_domain(model)

    def find_recent_active_by_library_and_name(
        self,
        library_id: str,
        name: str,
        within_minutes: int = 30,
        file_size: int = 0,
    ) -> Asset | None:
        """兜底去重：同库 + 同文件名（+同大小）且近期仍处活动状态（uploading/processing）的素材。

        用于旧客户端未传 file_hash/client_upload_id 时，防止 complete 超时重试
        反复创建 PROCESSING 占位记录。只命中"活动中"的近期记录，READY 历史素材不拦。

        严格模式（#1714 误杀修复）：file_size 必须 > 0 且与记录大小严格一致；
        file_size=0（大小未知）时直接返回 None——宁可漏判（极端情况下多建一条
        占位）也不可仅凭同名 + processing 误杀内容全新的视频。
        """
        from datetime import datetime, timedelta

        if not name:
            return None
        if not file_size or file_size <= 0:
            return None
        cutoff = datetime.now(UTC) - timedelta(minutes=within_minutes)
        query = self.session.query(AssetModel).filter(
            AssetModel.asset_library_id == library_id,
            AssetModel.name == name,
            AssetModel.status.in_([AssetStatus.UPLOADING.value, AssetStatus.PROCESSING.value]),
            AssetModel.created_at >= cutoff,
            AssetModel.file_size == file_size,
        )
        model = query.order_by(AssetModel.created_at.desc()).first()
        if model is None:
            return None
        return self._to_domain(model)