SupplyAgent/supply_infra/db/repositories/video_discovery_repo.py

from __future__ import annotations

import json
from typing import Any

from sqlalchemy import func, select
from sqlalchemy.exc import OperationalError, ProgrammingError

from supply_infra.video_discovery_gates import (
    evaluate_candidate_gate,
    load_rule_snapshot,
)
from supply_infra.db.models.video_discovery import (
    VideoDiscoveryCandidate,
    VideoDiscoveryRun,
    VideoDiscoverySearch,
)
from supply_infra.db.repositories.base import BaseRepository

def _json_list(raw: str | None) -> list[Any]:
    if not raw:
        return []
    try:
        value = json.loads(raw)
    except (TypeError, ValueError):
        return []
    return value if isinstance(value, list) else []


def _merge_json_list(raw: str | None, values: list[Any]) -> str | None:
    merged = _json_list(raw)
    seen = {
        json.dumps(value, ensure_ascii=False, sort_keys=True, default=str)
        for value in merged
    }
    for value in values:
        key = json.dumps(value, ensure_ascii=False, sort_keys=True, default=str)
        if key not in seen:
            seen.add(key)
            merged.append(value)
    return json.dumps(merged, ensure_ascii=False) if merged else None


_PUBLISHABLE_BUCKETS = ("primary",)


class VideoDiscoveryRepository(BaseRepository[VideoDiscoveryRun]):
    """需求找视频运行、搜索轨迹和候选快照的统一 Repository。"""

    model = VideoDiscoveryRun

    def get_run(self, run_id: str) -> VideoDiscoveryRun | None:
        stmt = select(VideoDiscoveryRun).where(VideoDiscoveryRun.run_id == run_id)
        return self.session.scalar(stmt)

    def get_by_biz_dt_and_grade(
        self, biz_dt: str, demand_grade_id: int
    ) -> VideoDiscoveryRun | None:
        stmt = select(VideoDiscoveryRun).where(
            VideoDiscoveryRun.biz_dt == biz_dt,
            VideoDiscoveryRun.demand_grade_id == int(demand_grade_id),
        )
        return self.session.scalar(stmt)

    def list_skip_grade_ids(self, biz_dt: str) -> set[int]:
        """返回指定业务日已有候选落库的 demand_grade.id（与 run_outcome 成功口径一致）。"""
        stmt = (
            select(VideoDiscoveryRun.demand_grade_id)
            .join(
                VideoDiscoveryCandidate,
                VideoDiscoveryCandidate.run_id == VideoDiscoveryRun.run_id,
            )
            .where(
                VideoDiscoveryRun.biz_dt == biz_dt,
                VideoDiscoveryRun.demand_grade_id.is_not(None),
            )
            .distinct()
        )
        try:
            return {
                int(grade_id)
                for grade_id in self.session.scalars(stmt).all()
                if grade_id is not None
            }
        except (OperationalError, ProgrammingError) as exc:
            if "biz_dt" not in str(exc).lower():
                raise
            return set()

    def upsert_scheduled_run(self, values: dict[str, Any]) -> VideoDiscoveryRun:
        """按 (biz_dt, demand_grade_id) 预创建或重置调度运行记录。"""
        biz_dt = str(values["biz_dt"])
        demand_grade_id = int(values["demand_grade_id"])
        existing = self.get_by_biz_dt_and_grade(biz_dt, demand_grade_id)
        if existing is None:
            entity = VideoDiscoveryRun(**values)
            self.session.add(entity)
            self.session.flush()
            return entity

        for key, value in values.items():
            if key in {"id", "create_time"}:
                continue
            if hasattr(existing, key):
                setattr(existing, key, value)
        existing.status = str(values.get("status") or "running")
        existing.stop_reason = values.get("stop_reason")
        self.session.flush()
        return existing

    def mark_run_failed(self, run_id: str, *, stop_reason: str | None = None) -> None:
        run = self.get_run(run_id)
        if run is None:
            return
        run.status = "failed"
        if stop_reason:
            run.stop_reason = stop_reason[:2000]
        self.session.flush()

    def create_run(self, values: dict[str, Any]) -> VideoDiscoveryRun:
        entity = VideoDiscoveryRun(**values)
        self.session.add(entity)
        self.session.flush()
        return entity

    def save_search_page(
        self,
        search_values: dict[str, Any],
        candidate_rows: list[dict[str, Any]],
    ) -> tuple[VideoDiscoverySearch, list[VideoDiscoveryCandidate]]:
        """新增一个搜索页，并为本页每条结果新增独立候选记录。"""
        run_id = str(search_values["run_id"])
        search = VideoDiscoverySearch(**search_values)
        self.session.add(search)
        self.session.flush()

        candidates: list[VideoDiscoveryCandidate] = []
        aweme_ids: list[str] = []
        for raw_row in candidate_rows:
            row = dict(raw_row)
            aweme_id = str(row.pop("aweme_id", "") or "").strip()
            if not aweme_id:
                continue
            source_keyword = row.pop("_source_keyword", None)
            entity = VideoDiscoveryCandidate(
                run_id=run_id,
                search_id=int(search.id),
                aweme_id=aweme_id,
                source_keywords_json=(
                    json.dumps([source_keyword], ensure_ascii=False)
                    if source_keyword
                    else None
                ),
                source_search_ids_json=json.dumps([int(search.id)]),
                decision_bucket="pending_evaluation",
            )
            for key, value in row.items():
                if value is None:
                    continue
                if hasattr(entity, key):
                    setattr(entity, key, value)
            self.session.add(entity)
            candidates.append(entity)
            aweme_ids.append(aweme_id)

        search.new_candidate_count = len(candidates)
        search.result_ids_json = (
            json.dumps(aweme_ids, ensure_ascii=False) if aweme_ids else None
        )
        self.session.flush()
        self._refresh_run_counts(run_id)
        return search, candidates

    def update_candidates(
        self,
        run_id: str,
        rows: list[dict[str, Any]],
    ) -> list[VideoDiscoveryCandidate]:
        """严格按 candidate_id 更新候选；不新增记录、不修改运行状态。"""
        candidate_ids = [int(row["candidate_id"]) for row in rows]
        if len(candidate_ids) != len(set(candidate_ids)):
            raise ValueError("candidate_id 不能重复")

        run = self.get_run(run_id)
        if run is None:
            raise ValueError(f"run_id 不存在: {run_id}")
        rule_snapshot = load_rule_snapshot(run.rule_config_json)

        stmt = select(VideoDiscoveryCandidate).where(
            VideoDiscoveryCandidate.run_id == run_id,
            VideoDiscoveryCandidate.id.in_(candidate_ids),
        )
        existing = {
            int(item.id): item for item in self.session.scalars(stmt).all()
        }
        missing = [candidate_id for candidate_id in candidate_ids if candidate_id not in existing]
        if missing:
            raise ValueError(
                f"candidate_id 不存在或不属于 run_id={run_id}: {missing}"
            )

        gate_failures: list[str] = []
        for row in rows:
            candidate_id = int(row["candidate_id"])
            entity = existing[candidate_id]
            for key, value in row.items():
                if key in {"candidate_id", "id", "run_id", "search_id", "aweme_id"}:
                    continue
                if value is None:
                    continue
                if key in {
                    "source_keywords_json",
                    "source_search_ids_json",
                    "tags_json",
                }:
                    values = value if isinstance(value, list) else _json_list(str(value))
                    setattr(entity, key, _merge_json_list(getattr(entity, key), values))
                elif hasattr(entity, key):
                    setattr(entity, key, value)

            gate = evaluate_candidate_gate(
                {
                    "title": entity.title,
                    "tags_json": entity.tags_json,
                    "publish_at": entity.publish_at,
                    "duration_seconds": entity.duration_seconds,
                    "share_count": entity.share_count,
                    "content_50_plus_ratio": entity.content_50_plus_ratio,
                    "account_50_plus_ratio": entity.account_50_plus_ratio,
                    "temporal_type": entity.temporal_type,
                    "temporal_status": entity.temporal_status,
                    "temporal_evidence_json": entity.temporal_evidence_json,
                },
                rule_snapshot,
            )
            entity.content_portrait_status = gate["content_portrait_status"]
            entity.account_portrait_status = gate["account_portrait_status"]
            entity.portrait_conflict = int(bool(gate["portrait_conflict"]))
            entity.temporal_type = gate["temporal"]["temporal_type"]
            entity.temporal_status = gate["temporal"]["status"]
            entity.temporal_evidence_json = json.dumps(
                gate["temporal"],
                ensure_ascii=False,
                default=str,
            )
            entity.gate_status = gate["status"]
            entity.gate_results_json = json.dumps(
                gate,
                ensure_ascii=False,
                default=str,
            )
            entity.rule_version = str(gate["rule_version"])

            if entity.decision_bucket == "primary":
                if not gate["primary_eligible"]:
                    codes = ", ".join(gate["failed_reason_codes"])
                    gate_failures.append(f"candidate_id={candidate_id}: {codes}")
                else:
                    entity.reject_reason_code = None
            else:
                failed_codes = gate["failed_reason_codes"]
                if failed_codes:
                    entity.reject_reason_code = str(failed_codes[0])
                elif not entity.reject_reason_code:
                    entity.reject_reason_code = "REJECTED_BY_AGENT"

        if gate_failures:
            raise ValueError(
                "primary 候选未通过 P0 硬门槛: " + "; ".join(gate_failures)
            )
        self.session.flush()
        return [existing[candidate_id] for candidate_id in candidate_ids]

    def finish_run(
        self,
        run_id: str,
        *,
        status: str,
        intent_summary: str | None = None,
        stop_reason: str | None = None,
    ) -> VideoDiscoveryRun:
        run = self.get_run(run_id)
        if run is None:
            raise ValueError(f"run_id 不存在: {run_id}")
        run.status = status
        if intent_summary is not None:
            run.intent_summary = intent_summary
        if stop_reason is not None:
            run.stop_reason = stop_reason
        self.session.flush()
        self._refresh_run_counts(run_id)
        return run

    def list_searches(self, run_id: str) -> list[VideoDiscoverySearch]:
        stmt = (
            select(VideoDiscoverySearch)
            .where(VideoDiscoverySearch.run_id == run_id)
            .order_by(VideoDiscoverySearch.id)
        )
        return list(self.session.scalars(stmt).all())

    def list_candidates(
        self,
        run_id: str,
        *,
        buckets: tuple[str, ...] | None = None,
        limit: int = 100,
    ) -> list[VideoDiscoveryCandidate]:
        stmt = select(VideoDiscoveryCandidate).where(
            VideoDiscoveryCandidate.run_id == run_id
        )
        if buckets:
            stmt = stmt.where(VideoDiscoveryCandidate.decision_bucket.in_(buckets))
        stmt = stmt.order_by(
            VideoDiscoveryCandidate.value_score.desc(),
            VideoDiscoveryCandidate.share_count.desc(),
            VideoDiscoveryCandidate.id,
        ).limit(limit)
        return list(self.session.scalars(stmt).all())

    def has_candidates(self, run_id: str) -> bool:
        stmt = (
            select(VideoDiscoveryCandidate.id)
            .where(VideoDiscoveryCandidate.run_id == run_id)
            .limit(1)
        )
        return self.session.scalar(stmt) is not None

    def list_publishable_candidates(
        self,
        *,
        run_id: str | None = None,
        biz_dt: str | None = None,
        skip_published: bool = True,
        limit: int | None = None,
    ) -> list[VideoDiscoveryCandidate]:
        """查询待提交 AIGC 的候选视频；仅 primary 分池，且 aweme_id 非空。"""
        stmt = (
            select(VideoDiscoveryCandidate)
            .join(
                VideoDiscoveryRun,
                VideoDiscoveryRun.run_id == VideoDiscoveryCandidate.run_id,
            )
            .where(VideoDiscoveryCandidate.decision_bucket.in_(_PUBLISHABLE_BUCKETS))
            .where(VideoDiscoveryCandidate.aweme_id.is_not(None))
            .where(func.trim(VideoDiscoveryCandidate.aweme_id) != "")
            .order_by(
                VideoDiscoveryCandidate.value_score.desc(),
                VideoDiscoveryCandidate.share_count.desc(),
                VideoDiscoveryCandidate.id,
            )
        )
        if run_id:
            stmt = stmt.where(VideoDiscoveryCandidate.run_id == run_id)
        if biz_dt:
            stmt = stmt.where(VideoDiscoveryRun.biz_dt == biz_dt)
        if skip_published:
            stmt = stmt.where(VideoDiscoveryCandidate.aigc_crawler_plan_id.is_(None))
        if limit is not None:
            stmt = stmt.limit(limit)
        return list(self.session.scalars(stmt).all())

    def count_passed_videos(self, biz_dt: str) -> int:
        """统计业务日内 primary 的唯一视频数。"""
        stmt = (
            select(func.count(func.distinct(VideoDiscoveryCandidate.aweme_id)))
            .join(
                VideoDiscoveryRun,
                VideoDiscoveryRun.run_id == VideoDiscoveryCandidate.run_id,
            )
            .where(VideoDiscoveryRun.biz_dt == biz_dt)
            .where(VideoDiscoveryCandidate.decision_bucket.in_(_PUBLISHABLE_BUCKETS))
            .where(VideoDiscoveryCandidate.aweme_id.is_not(None))
            .where(func.trim(VideoDiscoveryCandidate.aweme_id) != "")
        )
        return int(self.session.scalar(stmt) or 0)

    def mark_candidates_aigc_plans(
        self,
        candidate_ids: list[int],
        *,
        crawler_plan_id: str,
        produce_plan_id: str,
        publish_plan_id: str,
        plan_label: str,
    ) -> int:
        if not candidate_ids:
            return 0
        stmt = select(VideoDiscoveryCandidate).where(
            VideoDiscoveryCandidate.id.in_(candidate_ids)
        )
        updated = 0
        for entity in self.session.scalars(stmt).all():
            entity.aigc_crawler_plan_id = crawler_plan_id
            entity.aigc_produce_plan_id = produce_plan_id
            entity.aigc_publish_plan_id = publish_plan_id
            entity.aigc_plan_label = plan_label
            updated += 1
        self.session.flush()
        return updated

    def _refresh_run_counts(self, run_id: str) -> None:
        run = self.get_run(run_id)
        if run is None:
            return

        search_stmt = select(func.count(VideoDiscoverySearch.id)).where(
            VideoDiscoverySearch.run_id == run_id
        )
        run.search_count = int(self.session.scalar(search_stmt) or 0)

        bucket_stmt = (
            select(
                VideoDiscoveryCandidate.decision_bucket,
                func.count(VideoDiscoveryCandidate.id),
            )
            .where(VideoDiscoveryCandidate.run_id == run_id)
            .group_by(VideoDiscoveryCandidate.decision_bucket)
        )
        counts = {str(bucket): int(count) for bucket, count in self.session.execute(bucket_stmt)}
        run.primary_count = counts.get("primary", 0)
        self.session.flush()