SupplyAgent/supply_infra/db/repositories/multi_demand_video_detail_repo.py

from __future__ import annotations

from collections.abc import Iterable

from sqlalchemy import select, update
from sqlalchemy.dialects.mysql import insert

from supply_infra.db.models.multi_demand_video_detail import MultiDemandVideoDetail
from supply_infra.db.repositories.base import BaseRepository

_BATCH_SIZE = 1000


class MultiDemandVideoDetailRepository(BaseRepository[MultiDemandVideoDetail]):
    """需求池视频详情表 — 按 vid 唯一键增量写入。"""

    model = MultiDemandVideoDetail

    def get_existing_vids(self, vids: Iterable[str]) -> set[str]:
        """返回 vids 中已存在于表内的视频 id。"""
        vid_list = [v for v in vids if v]
        if not vid_list:
            return set()

        existing: set[str] = set()
        for i in range(0, len(vid_list), _BATCH_SIZE):
            batch = vid_list[i : i + _BATCH_SIZE]
            stmt = select(MultiDemandVideoDetail.vid).where(
                MultiDemandVideoDetail.vid.in_(batch)
            )
            existing.update(str(v) for v in self.session.scalars(stmt).all() if v)
        return existing

    def list_by_vids(self, vids: Iterable[str]) -> dict[str, MultiDemandVideoDetail]:
        """按 vid 批量查询，返回 vid → row。"""
        vid_list = [v for v in vids if v]
        if not vid_list:
            return {}

        result: dict[str, MultiDemandVideoDetail] = {}
        for i in range(0, len(vid_list), _BATCH_SIZE):
            batch = vid_list[i : i + _BATCH_SIZE]
            stmt = select(MultiDemandVideoDetail).where(
                MultiDemandVideoDetail.vid.in_(batch)
            )
            for row in self.session.scalars(stmt).all():
                if row.vid:
                    result[str(row.vid)] = row
        return result

    def list_vids_missing_title(self) -> list[str]:
        """返回 title 为空的全部 vid。"""
        stmt = select(MultiDemandVideoDetail.vid).where(
            (MultiDemandVideoDetail.title.is_(None)) | (MultiDemandVideoDetail.title == "")
        )
        return [str(v) for v in self.session.scalars(stmt).all() if v]

    def list_vids_missing_points(self) -> list[str]:
        """返回灵感点/目的点/关键点任一字段为空的全部 vid。"""
        stmt = select(MultiDemandVideoDetail.vid).where(
            MultiDemandVideoDetail.inspiration_points_json.is_(None)
            | MultiDemandVideoDetail.purpose_points_json.is_(None)
            | MultiDemandVideoDetail.key_points_json.is_(None)
        )
        return [str(v) for v in self.session.scalars(stmt).all() if v]

    def list_vids_missing_urls(self) -> list[str]:
        """返回 url1 或 url2 为空的全部 vid。"""
        stmt = select(MultiDemandVideoDetail.vid).where(
            MultiDemandVideoDetail.url1.is_(None)
            | (MultiDemandVideoDetail.url1 == "")
            | MultiDemandVideoDetail.url2.is_(None)
            | (MultiDemandVideoDetail.url2 == "")
        )
        return [str(v) for v in self.session.scalars(stmt).all() if v]

    def bulk_insert_ignore(self, rows: list[dict]) -> int:
        """批量插入，MySQL 自动忽略 vid 重复行。"""
        if not rows:
            return 0

        inserted = 0
        for i in range(0, len(rows), _BATCH_SIZE):
            batch = rows[i : i + _BATCH_SIZE]
            stmt = insert(MultiDemandVideoDetail).values(batch).prefix_with("IGNORE")
            result = self.session.execute(stmt)
            inserted += result.rowcount
        return inserted

    def update_titles(self, titles_by_vid: dict[str, str]) -> int:
        """按 vid 批量更新 title。"""
        if not titles_by_vid:
            return 0

        updated = 0
        for vid, title in titles_by_vid.items():
            stmt = (
                update(MultiDemandVideoDetail)
                .where(MultiDemandVideoDetail.vid == vid)
                .values(title=title)
            )
            result = self.session.execute(stmt)
            updated += result.rowcount or 0
        return updated

    def update_urls(self, urls_by_vid: dict[str, dict[str, str | None]]) -> int:
        """按 vid 批量更新 url1/url2（values 为落库字段名 → URL 文本）。"""
        if not urls_by_vid:
            return 0

        updated = 0
        for vid, values in urls_by_vid.items():
            if not values:
                continue
            stmt = (
                update(MultiDemandVideoDetail)
                .where(MultiDemandVideoDetail.vid == vid)
                .values(**values)
            )
            result = self.session.execute(stmt)
            updated += result.rowcount or 0
        return updated

    def update_points(self, points_by_vid: dict[str, dict[str, str | None]]) -> int:
        """按 vid 批量更新灵感点/目的点/关键点（values 为落库字段名 → JSON 文本）。"""
        if not points_by_vid:
            return 0

        updated = 0
        for vid, values in points_by_vid.items():
            if not values:
                continue
            stmt = (
                update(MultiDemandVideoDetail)
                .where(MultiDemandVideoDetail.vid == vid)
                .values(**values)
            )
            result = self.session.execute(stmt)
            updated += result.rowcount or 0
        return updated