from __future__ import annotations

import json
import hashlib
from dataclasses import dataclass
from datetime import datetime
from typing import Any

from sqlalchemy import func, select, update
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,
    VideoDiscoveryEvidence,
    VideoDiscoveryGateEvaluation,
    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",)
_BUSINESS_GOAL_PRIMARY_COUNT = 5
_PENDING_FINAL_REJECT_REASON = "NOT_SELECTED_AFTER_EVALUATION"
_EVIDENCE_INCOMPLETE_REJECT_REASON = "EVIDENCE_INCOMPLETE"


def _gate_input(entity: VideoDiscoveryCandidate) -> dict[str, Any]:
    """Build the single authoritative gate input from persisted candidate facts."""
    return {
        "title": entity.title,
        "tags_json": entity.tags_json,
        "publish_at": entity.publish_at,
        "duration_seconds": entity.duration_seconds,
        "share_count": entity.share_count,
        "like_count": entity.like_count,
        "play_count": entity.play_count,
        "content_50_plus_ratio": entity.content_50_plus_ratio,
        "account_50_plus_ratio": entity.account_50_plus_ratio,
        "relevance_score": entity.relevance_score,
        "elder_score": entity.elder_score,
        "share_score": entity.share_score,
        "value_score": entity.value_score,
        "temporal_type": entity.temporal_type,
        "temporal_status": entity.temporal_status,
        "temporal_evidence_json": entity.temporal_evidence_json,
        "detail_fetch_status": entity.detail_fetch_status,
        "content_portrait_fetch_status": entity.content_portrait_fetch_status,
        "account_portrait_fetch_status": entity.account_portrait_fetch_status,
        "latest_search_evidence_id": entity.latest_search_evidence_id,
        "latest_detail_evidence_id": entity.latest_detail_evidence_id,
        "latest_content_portrait_evidence_id": entity.latest_content_portrait_evidence_id,
        "latest_account_portrait_evidence_id": entity.latest_account_portrait_evidence_id,
        "evidence_version": int(entity.evidence_version or 0),
    }


def _pending_gate_codes(
    gate: dict[str, Any], candidate: VideoDiscoveryCandidate
) -> list[str]:
    """Return missing-evidence failures that are still retryable."""
    pending: list[str] = []
    detail_terminal = candidate.detail_fetch_status in {"success", "unavailable"}
    portrait_terminal = (
        candidate.content_portrait_fetch_status in {"success", "unavailable"}
        or candidate.account_portrait_fetch_status in {"success", "unavailable"}
    )
    for code in gate.get("failed_reason_codes") or []:
        code_text = str(code)
        if code_text in {
            "TEMPORAL_UNKNOWN",
            "DURATION_UNKNOWN",
            "SHARE_COUNT_UNKNOWN",
        } and not detail_terminal:
            pending.append(code_text)
        elif code_text == "CONTENT_PORTRAIT_MISSING" and not portrait_terminal:
            pending.append(code_text)
    return pending


@dataclass(frozen=True)
class CandidateUpdateBatch:
    candidates: list[VideoDiscoveryCandidate]
    reclassified: list[dict[str, Any]]


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。"""
        stmt = select(VideoDiscoveryRun.demand_grade_id).where(
            VideoDiscoveryRun.biz_dt == biz_dt,
            VideoDiscoveryRun.demand_grade_id.is_not(None),
            (
                (VideoDiscoveryRun.status == "finished")
                | (VideoDiscoveryRun.attempt_count >= 1)
            ),
        )
        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

        next_attempt = int(existing.attempt_count or 0) + 1
        for key, value in values.items():
            if key in {"id", "create_time"}:
                continue
            if key == "attempt_count":
                continue
            if hasattr(existing, key):
                setattr(existing, key, value)
        existing.status = str(values.get("status") or "running")
        existing.outcome_status = None
        existing.attempt_count = next_attempt
        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"
        run.outcome_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 create_evidence(self, values: dict[str, Any]) -> VideoDiscoveryEvidence:
        entity = VideoDiscoveryEvidence(**values)
        self.session.add(entity)
        self.session.flush()
        return entity

    def get_evidence(self, evidence_id: int) -> VideoDiscoveryEvidence | None:
        return self.session.get(VideoDiscoveryEvidence, int(evidence_id))

    def get_candidates_by_ids(
        self, run_id: str, candidate_ids: list[int]
    ) -> list[VideoDiscoveryCandidate]:
        if not candidate_ids:
            return []
        stmt = select(VideoDiscoveryCandidate).where(
            VideoDiscoveryCandidate.run_id == run_id,
            VideoDiscoveryCandidate.id.in_([int(value) for value in candidate_ids]),
        )
        rows = list(self.session.scalars(stmt).all())
        by_id = {int(item.id): item for item in rows}
        missing = [int(value) for value in candidate_ids if int(value) not in by_id]
        if missing:
            raise ValueError(
                f"candidate_id 不存在或不属于 run_id={run_id}: {missing}"
            )
        return [by_id[int(value)] for value in candidate_ids]

    def _evaluate_candidate(
        self,
        entity: VideoDiscoveryCandidate,
        rule_snapshot: dict[str, Any],
    ) -> dict[str, Any]:
        input_payload = _gate_input(entity)
        gate = evaluate_candidate_gate(input_payload, rule_snapshot)
        pending_codes = _pending_gate_codes(gate, entity)
        if pending_codes:
            gate["status"] = "pending"
            gate["primary_eligible"] = False
            gate["pending_reason_codes"] = pending_codes

        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 = str(gate["status"])
        entity.gate_results_json = json.dumps(gate, ensure_ascii=False, default=str)
        entity.rule_version = str(gate["rule_version"])
        entity.gate_evidence_version = int(entity.evidence_version or 0)
        entity.gate_evaluated_at = datetime.now()

        input_json = json.dumps(
            input_payload, ensure_ascii=False, sort_keys=True, default=str
        )
        input_hash = hashlib.sha256(input_json.encode("utf-8")).hexdigest()
        evaluation = VideoDiscoveryGateEvaluation(
            run_id=entity.run_id,
            candidate_id=int(entity.id),
            rule_version=str(gate["rule_version"]),
            evidence_version=int(entity.evidence_version or 0),
            gate_input_hash=input_hash,
            gate_input_json=input_json,
            gate_result_json=json.dumps(gate, ensure_ascii=False, default=str),
            gate_status=str(gate["status"]),
            failed_reason_codes_json=json.dumps(
                gate.get("failed_reason_codes") or [], ensure_ascii=False
            ),
        )
        self.session.add(evaluation)
        self.session.flush()
        entity.latest_gate_evaluation_id = int(evaluation.id)
        return gate

    def apply_evidence(
        self,
        run_id: str,
        *,
        aweme_id: str,
        evidence: VideoDiscoveryEvidence,
        normalized: dict[str, Any],
    ) -> list[VideoDiscoveryCandidate]:
        """Apply one immutable observation to every occurrence of a run video."""
        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.aweme_id == aweme_id,
        )
        candidates = list(self.session.scalars(stmt).all())
        if not candidates:
            raise ValueError(f"run_id={run_id} 中不存在 aweme_id={aweme_id}")

        pointer_name = {
            "detail": "latest_detail_evidence_id",
            "content_portrait": "latest_content_portrait_evidence_id",
            "account_portrait": "latest_account_portrait_evidence_id",
        }.get(evidence.evidence_type)
        status_name = {
            "detail": "detail_fetch_status",
            "content_portrait": "content_portrait_fetch_status",
            "account_portrait": "account_portrait_fetch_status",
        }.get(evidence.evidence_type)
        for candidate in candidates:
            if pointer_name:
                setattr(candidate, pointer_name, int(evidence.id))
            if status_name:
                setattr(candidate, status_name, evidence.fetch_status)
            for key, value in normalized.items():
                if key in {"id", "candidate_id", "run_id", "search_id", "aweme_id"}:
                    continue
                if value is not None and hasattr(candidate, key):
                    setattr(candidate, key, value)
            candidate.evidence_version = int(candidate.evidence_version or 0) + 1
            self.session.flush()
            gate = self._evaluate_candidate(candidate, rule_snapshot)
            if candidate.decision_bucket == "primary" and gate["status"] != "pass":
                if gate["status"] == "pending":
                    candidate.decision_bucket = "pending_evaluation"
                    candidate.reject_reason_code = None
                else:
                    candidate.decision_bucket = "rejected"
                    failed_codes = gate.get("failed_reason_codes") or []
                    candidate.reject_reason_code = str(
                        failed_codes[0] if failed_codes else "P0_GATE_FAILED"
                    )
        evidence.processing_status = "applied"
        evidence.normalized_json = json.dumps(
            normalized, ensure_ascii=False, default=str
        )
        self.session.flush()
        self._refresh_run_counts(run_id)
        return candidates

    def save_search_page(
        self,
        search_values: dict[str, Any],
        candidate_rows: list[dict[str, Any]],
        evidence_values: dict[str, Any] | None = None,
        evidence_id: int | None = None,
    ) -> tuple[VideoDiscoverySearch, list[VideoDiscoveryCandidate]]:
        """新增一个搜索页，并为本页每条结果新增独立候选记录。"""
        run_id = str(search_values["run_id"])
        evidence = (
            self.get_evidence(evidence_id)
            if evidence_id is not None
            else (self.create_evidence(evidence_values) if evidence_values else None)
        )
        if evidence_id is not None and evidence is None:
            raise ValueError(f"evidence_id 不存在: {evidence_id}")
        search = VideoDiscoverySearch(**search_values)
        if evidence is not None:
            search.raw_evidence_id = int(evidence.id)
        self.session.add(search)
        self.session.flush()
        if evidence is not None:
            evidence.search_id = int(search.id)

        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",
                latest_search_evidence_id=(int(evidence.id) if evidence else None),
                evidence_version=(1 if evidence else 0),
            )
            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()
        if evidence is not None:
            evidence.processing_status = "applied"
        self._refresh_run_counts(run_id)
        return search, candidates

    def update_candidates(
        self,
        run_id: str,
        rows: list[dict[str, Any]],
    ) -> CandidateUpdateBatch:
        """严格按 candidate_id 更新候选；P0 不合格 primary 自动降级为 rejected。"""
        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}"
            )

        reclassified: list[dict[str, Any]] = []
        for row in rows:
            candidate_id = int(row["candidate_id"])
            entity = existing[candidate_id]
            requested_bucket = str(row.get("decision_bucket") or entity.decision_bucket)
            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)
            entity.evidence_version = int(entity.evidence_version or 0) + 1
            self.session.flush()
            gate = self._evaluate_candidate(entity, rule_snapshot)

            if requested_bucket == "primary":
                if gate["status"] == "pending":
                    entity.decision_bucket = "pending_evaluation"
                    entity.reject_reason_code = None
                    reclassified.append(
                        {
                            "candidate_id": candidate_id,
                            "requested_bucket": "primary",
                            "saved_bucket": "pending_evaluation",
                            "failed_reason_codes": list(
                                gate.get("pending_reason_codes") or []
                            ),
                        }
                    )
                elif not gate["primary_eligible"]:
                    failed_codes = gate["failed_reason_codes"]
                    entity.decision_bucket = "rejected"
                    entity.reject_reason_code = str(
                        failed_codes[0] if failed_codes else "P0_GATE_FAILED"
                    )
                    reclassified.append(
                        {
                            "candidate_id": candidate_id,
                            "requested_bucket": "primary",
                            "saved_bucket": "rejected",
                            "failed_reason_codes": list(failed_codes),
                        }
                    )
                else:
                    entity.decision_bucket = "primary"
                    entity.reject_reason_code = None
            else:
                entity.decision_bucket = requested_bucket
                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"

        self.session.flush()
        self._refresh_run_counts(run_id)
        return CandidateUpdateBatch(
            candidates=[existing[candidate_id] for candidate_id in candidate_ids],
            reclassified=reclassified,
        )

    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 status == "finished":
            self._finalize_pending_candidates(run_id)
        elif status == "failed":
            run.outcome_status = "failed"
        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)
        if status == "finished":
            run.outcome_status = self._business_outcome(run.valid_primary_count)
        return run

    @staticmethod
    def _business_outcome(valid_primary_count: int) -> str:
        if valid_primary_count >= _BUSINESS_GOAL_PRIMARY_COUNT:
            return "goal_met"
        if valid_primary_count > 0:
            return "partial"
        return "no_match"

    def _finalize_pending_candidates(self, run_id: str) -> None:
        pending_buckets = ("pending_evaluation", "unreviewed")
        self.session.execute(
            update(VideoDiscoveryCandidate)
            .where(
                VideoDiscoveryCandidate.run_id == run_id,
                VideoDiscoveryCandidate.decision_bucket.in_(pending_buckets),
                VideoDiscoveryCandidate.gate_status == "pending",
            )
            .values(
                decision_bucket="rejected",
                reject_reason_code=_EVIDENCE_INCOMPLETE_REJECT_REASON,
            )
        )
        self.session.execute(
            update(VideoDiscoveryCandidate)
            .where(
                VideoDiscoveryCandidate.run_id == run_id,
                VideoDiscoveryCandidate.decision_bucket.in_(pending_buckets),
            )
            .values(
                decision_bucket="rejected",
                reject_reason_code=_PENDING_FINAL_REJECT_REASON,
            )
        )
        self.session.flush()

    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.gate_status == "pass")
            .where(
                VideoDiscoveryCandidate.gate_evidence_version
                == VideoDiscoveryCandidate.evidence_version
            )
            .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.gate_status == "pass")
            .where(
                VideoDiscoveryCandidate.gate_evidence_version
                == VideoDiscoveryCandidate.evidence_version
            )
            .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)
        valid_primary_stmt = select(
            func.count(func.distinct(VideoDiscoveryCandidate.aweme_id))
        ).where(
            VideoDiscoveryCandidate.run_id == run_id,
            VideoDiscoveryCandidate.decision_bucket == "primary",
            VideoDiscoveryCandidate.gate_status == "pass",
            VideoDiscoveryCandidate.gate_evidence_version
            == VideoDiscoveryCandidate.evidence_version,
            VideoDiscoveryCandidate.aweme_id.is_not(None),
            func.trim(VideoDiscoveryCandidate.aweme_id) != "",
        )
        run.valid_primary_count = int(
            self.session.scalar(valid_primary_stmt) or 0
        )
        self.session.flush()
