SupplyAgent/supply_infra/db/repositories/multi_demand_pool_di_repo.py

from __future__ import annotations

from sqlalchemy import delete, func, select, tuple_, update
from sqlalchemy.dialects.mysql import insert

from supply_infra.db.models.multi_demand_pool_di import MultiDemandPoolDi
from supply_infra.db.repositories.base import BaseRepository

_BATCH_SIZE = 1000

RowKey = tuple[str, str]
RealMetric = tuple[float | None, float | None]


class MultiDemandPoolDiRepository(BaseRepository[MultiDemandPoolDi]):
    """策略需求天级表 repository。"""

    model = MultiDemandPoolDi

    def get_latest_biz_dt(self) -> str | None:
        """返回需求池中最新业务日；无数据时返回 None。"""
        stmt = select(func.max(MultiDemandPoolDi.biz_dt))
        return self.session.scalar(stmt)

    def list_distinct_demand_name_summaries(
        self,
        biz_dt: str,
        *,
        limit: int = 50,
        offset: int = 0,
        exclude_names: list[str] | None = None,
    ) -> list[dict]:
        """
        按业务日分页列出去重需求词及聚合统计。

        返回每项包含：demand_name、row_count（出现行数）、strategies（策略列表）、
        total_weight、total_video_count、max_real_rov_7d、max_real_vov_7d。
        按 max_real_rov_7d 降序、total_weight 降序排列，优先暴露有后验数据/高权重的需求词。
        """
        stmt = (
            select(
                MultiDemandPoolDi.demand_name,
                func.count().label("row_count"),
                func.group_concat(MultiDemandPoolDi.strategy.distinct()).label("strategies"),
                func.sum(MultiDemandPoolDi.weight).label("total_weight"),
                func.sum(MultiDemandPoolDi.video_count).label("total_video_count"),
                func.max(MultiDemandPoolDi.real_rov_7d).label("max_real_rov_7d"),
                func.max(MultiDemandPoolDi.real_vov_7d).label("max_real_vov_7d"),
            )
            .where(MultiDemandPoolDi.biz_dt == biz_dt)
        )
        if exclude_names:
            stmt = stmt.where(MultiDemandPoolDi.demand_name.notin_(exclude_names))
        stmt = (
            stmt.group_by(MultiDemandPoolDi.demand_name)
            .order_by(
                func.max(MultiDemandPoolDi.real_rov_7d).desc(),
                func.sum(MultiDemandPoolDi.weight).desc(),
            )
            .limit(limit)
            .offset(offset)
        )

        rows = self.session.execute(stmt).all()
        return [
            {
                "demand_name": name,
                "row_count": int(row_count or 0),
                "strategies": (strategies or "").split(",") if strategies else [],
                "total_weight": float(total_weight) if total_weight is not None else None,
                "total_video_count": int(total_video_count) if total_video_count is not None else None,
                "max_real_rov_7d": float(max_rov) if max_rov is not None else None,
                "max_real_vov_7d": float(max_vov) if max_vov is not None else None,
            }
            for name, row_count, strategies, total_weight, total_video_count, max_rov, max_vov in rows
        ]

    def count_distinct_demand_names(self, biz_dt: str) -> int:
        """统计指定业务日去重需求词总数。"""
        stmt = select(func.count(func.distinct(MultiDemandPoolDi.demand_name))).where(
            MultiDemandPoolDi.biz_dt == biz_dt
        )
        return int(self.session.scalar(stmt) or 0)

    def list_by_biz_dt(self, biz_dt: str) -> list[MultiDemandPoolDi]:
        """返回业务日全部需求池行，供来源内排名等全局计算使用。"""
        stmt = (
            select(MultiDemandPoolDi)
            .where(MultiDemandPoolDi.biz_dt == biz_dt)
            .order_by(MultiDemandPoolDi.strategy, MultiDemandPoolDi.demand_name, MultiDemandPoolDi.id)
        )
        return list(self.session.scalars(stmt).all())

    def search_rows_by_name_fragment(self, biz_dt: str, keyword: str) -> list[dict]:
        """
        按业务日 + 需求名双向包含关系搜索明细行。

        匹配 demand_name LIKE %keyword% 或 keyword LIKE %demand_name%（互相包含），
        用于把同语义、措辞不同的需求词合并到一起判断。返回按 demand_name 去重后的明细。
        """
        keyword = (keyword or "").strip()
        if not keyword:
            return []

        stmt = select(MultiDemandPoolDi).where(
            MultiDemandPoolDi.biz_dt == biz_dt,
            MultiDemandPoolDi.demand_name.like(f"%{keyword}%"),
        )
        rows = list(self.session.scalars(stmt).all())

        if len(keyword) >= 2:
            broad_stmt = select(MultiDemandPoolDi).where(
                MultiDemandPoolDi.biz_dt == biz_dt,
            )
            seen_ids = {int(r.id) for r in rows}
            for row in self.session.scalars(broad_stmt).all():
                if int(row.id) in seen_ids:
                    continue
                name = row.demand_name or ""
                if name and name in keyword:
                    rows.append(row)
                    seen_ids.add(int(row.id))

        return [
            {
                "id": int(row.id),
                "demand_name": row.demand_name,
                "strategy": row.strategy,
                "weight": None if row.weight is None or row.weight == 0 else float(row.weight),
                "video_count": int(row.video_count) if row.video_count is not None else None,
                "real_rov_7d": float(row.real_rov_7d) if row.real_rov_7d is not None else None,
                "real_vov_7d": float(row.real_vov_7d) if row.real_vov_7d is not None else None,
            }
            for row in rows
        ]

    def get_by_ids(self, ids: list[int]) -> list[MultiDemandPoolDi]:
        """按 id 批量查询完整行。"""
        if not ids:
            return []
        rows: list[MultiDemandPoolDi] = []
        for start in range(0, len(ids), _BATCH_SIZE):
            batch = ids[start : start + _BATCH_SIZE]
            stmt = select(MultiDemandPoolDi).where(MultiDemandPoolDi.id.in_(batch))
            rows.extend(self.session.scalars(stmt).all())
        return rows

    def count_by_biz_dt(self, biz_dt: str) -> int:
        """统计指定业务日期去重行数（strategy + demand_id）。"""
        stmt = (
            select(func.count())
            .select_from(
                select(MultiDemandPoolDi.strategy, MultiDemandPoolDi.demand_id)
                .where(MultiDemandPoolDi.biz_dt == biz_dt)
                .distinct()
                .subquery()
            )
        )
        return int(self.session.scalar(stmt) or 0)

    def list_keys_by_biz_dt(self, biz_dt: str) -> set[RowKey]:
        """查询指定业务日期的 (strategy, demand_id) 键集合。"""
        stmt = select(MultiDemandPoolDi.strategy, MultiDemandPoolDi.demand_id).where(
            MultiDemandPoolDi.biz_dt == biz_dt
        )
        return {(str(s), str(d)) for s, d in self.session.execute(stmt).all()}

    def list_source_rows_by_biz_dt(self, biz_dt: str) -> list[dict]:
        """返回同步差分需要的源字段，不包含后续步骤回填的 ROV/VOV。"""
        stmt = select(
            MultiDemandPoolDi.id,
            MultiDemandPoolDi.strategy,
            MultiDemandPoolDi.demand_id,
            MultiDemandPoolDi.demand_name,
            MultiDemandPoolDi.weight,
            MultiDemandPoolDi.type,
            MultiDemandPoolDi.video_count,
            MultiDemandPoolDi.video_list,
            MultiDemandPoolDi.extend,
            MultiDemandPoolDi.biz_dt,
        ).where(MultiDemandPoolDi.biz_dt == biz_dt)
        return [
            {
                "id": int(row_id),
                "strategy": str(strategy),
                "demand_id": str(demand_id),
                "demand_name": str(demand_name),
                "weight": float(weight) if weight is not None else None,
                "type": str(demand_type) if demand_type is not None else None,
                "video_count": int(video_count) if video_count is not None else 0,
                "video_list": video_list,
                "extend": extend,
                "biz_dt": str(row_biz_dt),
            }
            for (
                row_id,
                strategy,
                demand_id,
                demand_name,
                weight,
                demand_type,
                video_count,
                video_list,
                extend,
                row_biz_dt,
            ) in self.session.execute(stmt).all()
        ]

    def delete_by_biz_dt(self, biz_dt: str) -> int:
        """删除指定业务日期的全部数据，便于按日重跑。"""
        stmt = delete(MultiDemandPoolDi).where(MultiDemandPoolDi.biz_dt == biz_dt)
        result = self.session.execute(stmt)
        return result.rowcount or 0

    def delete_by_keys(self, biz_dt: str, keys: list[RowKey]) -> int:
        """按 (strategy, demand_id) 批量删除指定业务日期数据。"""
        if not keys:
            return 0

        deleted = 0
        for i in range(0, len(keys), _BATCH_SIZE):
            batch = keys[i : i + _BATCH_SIZE]
            stmt = delete(MultiDemandPoolDi).where(
                MultiDemandPoolDi.biz_dt == biz_dt,
                tuple_(MultiDemandPoolDi.strategy, MultiDemandPoolDi.demand_id).in_(batch),
            )
            result = self.session.execute(stmt)
            deleted += result.rowcount or 0
        return deleted

    def list_demand_names_by_biz_dt(self, biz_dt: str) -> list[str]:
        """查询指定业务日期的全部 demand_name。"""
        stmt = select(MultiDemandPoolDi.demand_name).where(
            MultiDemandPoolDi.biz_dt == biz_dt
        )
        return [name for name in self.session.scalars(stmt).all() if name]

    def list_all_video_lists(self) -> list[str]:
        """查询全表非空 video_list（JSON 文本）。"""
        stmt = select(MultiDemandPoolDi.video_list).where(
            MultiDemandPoolDi.video_list.isnot(None),
            MultiDemandPoolDi.video_list != "",
        )
        return [text for text in self.session.scalars(stmt).all() if text]

    def list_id_name_video_lists(
        self,
        biz_dt: str,
    ) -> list[tuple[int, str, str | None]]:
        """查询指定业务日的 (id, demand_name, video_list)。"""
        stmt = select(
            MultiDemandPoolDi.id,
            MultiDemandPoolDi.demand_name,
            MultiDemandPoolDi.video_list,
        ).where(MultiDemandPoolDi.biz_dt == biz_dt)
        return [
            (int(row_id), str(name), video_list)
            for row_id, name, video_list in self.session.execute(stmt).all()
            if name is not None
        ]

    def list_weights_by_name_like(
        self,
        biz_dt: str,
        keyword: str,
        strategies: list[str],
    ) -> list[tuple[str, float | None]]:
        """按日期 + demand_name LIKE + 指定策略，返回 (strategy, weight)。"""
        if not keyword or not strategies:
            return []

        stmt = select(MultiDemandPoolDi.strategy, MultiDemandPoolDi.weight).where(
            MultiDemandPoolDi.biz_dt == biz_dt,
            MultiDemandPoolDi.demand_name.like(f"%{keyword}%"),
            MultiDemandPoolDi.strategy.in_(strategies),
        )
        return [
            (str(strategy), None if weight is None or weight == 0 else float(weight))
            for strategy, weight in self.session.execute(stmt).all()
        ]

    def list_real_metrics_by_name_like(
        self,
        biz_dt: str,
        keyword: str,
    ) -> list[tuple[float | None, float | None]]:
        """
        按日期 + demand_name LIKE 返回 rov_diff / vov_diff（real_rov_7d / real_vov_7d 字段）。

        同一 demand_name 只取一条（MAX），避免多策略重复膨胀 count。
        """
        if not keyword:
            return []

        stmt = (
            select(
                func.max(MultiDemandPoolDi.real_rov_7d),
                func.max(MultiDemandPoolDi.real_vov_7d),
            )
            .where(
                MultiDemandPoolDi.biz_dt == biz_dt,
                MultiDemandPoolDi.demand_name.like(f"%{keyword}%"),
            )
            .group_by(MultiDemandPoolDi.demand_name)
        )
        return [
            (
                float(rov) if rov is not None else None,
                float(vov) if vov is not None else None,
            )
            for rov, vov in self.session.execute(stmt).all()
        ]

    def bulk_insert(self, rows: list[dict]) -> int:
        """批量插入。"""
        if not rows:
            return 0

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

    def update_real_metrics_by_demand_name(
        self,
        biz_dt: str,
        metrics_by_name: dict[str, RealMetric],
    ) -> int:
        """按 demand_name（特征值）批量更新 rov_diff / vov_diff（real_rov_7d / real_vov_7d）。"""
        if not metrics_by_name:
            return 0

        updated = 0
        for demand_name, (real_rov_7d, real_vov_7d) in metrics_by_name.items():
            stmt = (
                update(MultiDemandPoolDi)
                .where(
                    MultiDemandPoolDi.biz_dt == biz_dt,
                    MultiDemandPoolDi.demand_name == demand_name,
                )
                .values(real_rov_7d=real_rov_7d, real_vov_7d=real_vov_7d)
            )
            result = self.session.execute(stmt)
            updated += result.rowcount or 0
        return updated

    def update_source_fields(self, biz_dt: str, rows: list[dict]) -> int:
        """按 (strategy, demand_id) 更新所有会影响下游的上游源字段。"""
        if not rows:
            return 0

        updated = 0
        for row in rows:
            stmt = (
                update(MultiDemandPoolDi)
                .where(
                    MultiDemandPoolDi.biz_dt == biz_dt,
                    MultiDemandPoolDi.strategy == row["strategy"],
                    MultiDemandPoolDi.demand_id == row["demand_id"],
                )
                .values(
                    demand_name=row["demand_name"],
                    video_list=row.get("video_list"),
                    video_count=row.get("video_count"),
                    weight=row.get("weight"),
                    type=row.get("type"),
                    extend=row.get("extend"),
                )
            )
            result = self.session.execute(stmt)
            updated += result.rowcount or 0
        return updated